class SpeciesArgsBuilder(ArgsBuilder): argnames = 'eps, nc, mass, zsp, density, temperature, forces' def __init__(self, species): super(SpeciesArgsBuilder, self).__init__() self._species = CompositeSpecies(species) def build(self): zsp = self._species.get_zsp('density') nc = [len(x) for x in zsp] self.args.update({'mass': self._species.get_mass()}) self.args.update({'nc': nc}) self.args.update({'zsp': np.array(np.broadcast_arrays(*zsp))}) self.args.update(self._build_density_temperature()) self.args.update(self._build_forces()) def _build_density_temperature(self): d = {} for p in ['density', 'temperature', 'eps']: profiles = self._species.get_profile(p) profiles = np.array(np.broadcast_arrays(*profiles)) d[p] = myrollaxis(profiles, 3) d['temperature'] = self.remove_dim_charge_state(d['temperature']) d['eps'] = self.remove_dim_charge_state(d['eps'])[..., 0] return d def _build_forces(self): gradients = [] for s in self._species: forces = self.get_forces(s) pressure_gradient = forces.get_pressure_gradient() temperature_gradient = forces.get_temperature_gradient() gradients.append( np.concatenate((pressure_gradient[..., None, :], temperature_gradient[..., None, :]), axis=-2)) forces = np.array(np.broadcast_arrays(*gradients)) return {'forces': myrollaxis(forces, 4)} def get_forces(self, s): return Forces(s.profiles, squeeze=False) def remove_dim_charge_state(self, a): if a.ndim == 3: return a[..., 0] else: return a