1import numpy as np 2 3from yt.fields.field_info_container import FieldInfoContainer 4from yt.fields.magnetic_field import setup_magnetic_field_aliases 5from yt.frontends.open_pmd.misc import is_const_component, parse_unit_dimension 6from yt.units.yt_array import YTQuantity 7from yt.utilities.logger import ytLogger as mylog 8from yt.utilities.on_demand_imports import _h5py as h5py 9from yt.utilities.physical_constants import mu_0, speed_of_light 10 11 12def setup_poynting_vector(self): 13 def _get_poyn(axis): 14 def poynting(field, data): 15 u = mu_0 ** -1 16 if axis in "x": 17 return u * ( 18 data[("openPMD", "E_y")] * data[("gas", "magnetic_field_z")] 19 - data[("openPMD", "E_z")] * data[("gas", "magnetic_field_y")] 20 ) 21 elif axis in "y": 22 return u * ( 23 data[("openPMD", "E_z")] * data[("gas", "magnetic_field_x")] 24 - data[("openPMD", "E_x")] * data[("gas", "magnetic_field_z")] 25 ) 26 elif axis in "z": 27 return u * ( 28 data[("openPMD", "E_x")] * data[("gas", "magnetic_field_y")] 29 - data[("openPMD", "E_y")] * data[("gas", "magnetic_field_x")] 30 ) 31 32 return poynting 33 34 for ax in "xyz": 35 self.add_field( 36 ("openPMD", f"poynting_vector_{ax}"), 37 sampling_type="cell", 38 function=_get_poyn(ax), 39 units="W/m**2", 40 ) 41 42 43def setup_kinetic_energy(self, ptype): 44 def _kin_en(field, data): 45 p2 = ( 46 data[ptype, "particle_momentum_x"] ** 2 47 + data[ptype, "particle_momentum_y"] ** 2 48 + data[ptype, "particle_momentum_z"] ** 2 49 ) 50 mass = data[ptype, "particle_mass"] * data[ptype, "particle_weighting"] 51 return ( 52 speed_of_light * np.sqrt(p2 + mass ** 2 * speed_of_light ** 2) 53 - mass * speed_of_light ** 2 54 ) 55 56 self.add_field( 57 (ptype, "particle_kinetic_energy"), 58 sampling_type="particle", 59 function=_kin_en, 60 units="kg*m**2/s**2", 61 ) 62 63 64def setup_velocity(self, ptype): 65 def _get_vel(axis): 66 def velocity(field, data): 67 c = speed_of_light 68 momentum = data[ptype, f"particle_momentum_{axis}"] 69 mass = data[ptype, "particle_mass"] 70 weighting = data[ptype, "particle_weighting"] 71 return momentum / np.sqrt( 72 (mass * weighting) ** 2 + (momentum ** 2) / (c ** 2) 73 ) 74 75 return velocity 76 77 for ax in "xyz": 78 self.add_field( 79 (ptype, f"particle_velocity_{ax}"), 80 sampling_type="particle", 81 function=_get_vel(ax), 82 units="m/s", 83 ) 84 85 86def setup_absolute_positions(self, ptype): 87 def _abs_pos(axis): 88 def ap(field, data): 89 return np.add( 90 data[ptype, f"particle_positionCoarse_{axis}"], 91 data[ptype, f"particle_positionOffset_{axis}"], 92 ) 93 94 return ap 95 96 for ax in "xyz": 97 self.add_field( 98 (ptype, f"particle_position_{ax}"), 99 sampling_type="particle", 100 function=_abs_pos(ax), 101 units="m", 102 ) 103 104 105class OpenPMDFieldInfo(FieldInfoContainer): 106 r"""Specifies which fields from the dataset yt should know about. 107 108 ``self.known_other_fields`` and ``self.known_particle_fields`` must be populated. 109 Entries for both of these lists must be tuples of the form ("name", ("units", 110 ["fields", "to", "alias"], "display_name")) These fields will be represented and 111 handled in yt in the way you define them here. The fields defined in both 112 ``self.known_other_fields`` and ``self.known_particle_fields`` will only be added to 113 a dataset (with units, aliases, etc), if they match any entry in the 114 ``OpenPMDHierarchy``'s ``self.field_list``. 115 116 Notes 117 ----- 118 119 Contrary to many other frontends, we dynamically obtain the known fields from the 120 simulation output. The openPMD markup is extremely flexible - names, dimensions and 121 the number of individual datasets can (and very likely will) vary. 122 123 openPMD states that record names and their components are only allowed to contain 124 * characters a-Z, 125 * the numbers 0-9 126 * and the underscore _ 127 * (equivalently, the regex \w). 128 Since yt widely uses the underscore in field names, openPMD's underscores (_) are 129 replaced by hyphen (-). 130 131 Derived fields will automatically be set up, if names and units of your known 132 on-disk (or manually derived) fields match the ones in [1]. 133 134 References 135 ---------- 136 * http://yt-project.org/docs/dev/analyzing/fields.html 137 * http://yt-project.org/docs/dev/developing/creating_frontend.html#data-meaning-structures 138 * https://github.com/openPMD/openPMD-standard/blob/latest/STANDARD.md 139 * [1] http://yt-project.org/docs/dev/reference/field_list.html#universal-fields 140 """ 141 142 _mag_fields = [] 143 144 def __init__(self, ds, field_list): 145 f = ds._handle 146 bp = ds.base_path 147 mp = ds.meshes_path 148 pp = ds.particles_path 149 150 try: 151 fields = f[bp + mp] 152 for fname in fields.keys(): 153 field = fields[fname] 154 if isinstance(field, h5py.Dataset) or is_const_component(field): 155 # Don't consider axes. 156 # This appears to be a vector field of single dimensionality 157 ytname = str("_".join([fname.replace("_", "-")])) 158 parsed = parse_unit_dimension( 159 np.asarray(field.attrs["unitDimension"], dtype="int64") 160 ) 161 unit = str(YTQuantity(1, parsed).units) 162 aliases = [] 163 # Save a list of magnetic fields for aliasing later on 164 # We can not reasonably infer field type/unit by name in openPMD 165 if unit == "T" or unit == "kg/(A*s**2)": 166 self._mag_fields.append(ytname) 167 self.known_other_fields += ((ytname, (unit, aliases, None)),) 168 else: 169 for axis in field.keys(): 170 ytname = str("_".join([fname.replace("_", "-"), axis])) 171 parsed = parse_unit_dimension( 172 np.asarray(field.attrs["unitDimension"], dtype="int64") 173 ) 174 unit = str(YTQuantity(1, parsed).units) 175 aliases = [] 176 # Save a list of magnetic fields for aliasing later on 177 # We can not reasonably infer field type by name in openPMD 178 if unit == "T" or unit == "kg/(A*s**2)": 179 self._mag_fields.append(ytname) 180 self.known_other_fields += ((ytname, (unit, aliases, None)),) 181 for i in self.known_other_fields: 182 mylog.debug("open_pmd - known_other_fields - %s", i) 183 except (KeyError, TypeError, AttributeError): 184 pass 185 186 try: 187 particles = f[bp + pp] 188 for pname in particles.keys(): 189 species = particles[pname] 190 for recname in species.keys(): 191 try: 192 record = species[recname] 193 parsed = parse_unit_dimension(record.attrs["unitDimension"]) 194 unit = str(YTQuantity(1, parsed).units) 195 ytattrib = str(recname).replace("_", "-") 196 if ytattrib == "position": 197 # Symbolically rename position to preserve yt's 198 # interpretation of the pfield particle_position is later 199 # derived in setup_absolute_positions in the way yt expects 200 ytattrib = "positionCoarse" 201 if isinstance(record, h5py.Dataset) or is_const_component( 202 record 203 ): 204 name = ["particle", ytattrib] 205 self.known_particle_fields += ( 206 (str("_".join(name)), (unit, [], None)), 207 ) 208 else: 209 for axis in record.keys(): 210 aliases = [] 211 name = ["particle", ytattrib, axis] 212 ytname = str("_".join(name)) 213 self.known_particle_fields += ( 214 (ytname, (unit, aliases, None)), 215 ) 216 except (KeyError): 217 if recname != "particlePatches": 218 mylog.info( 219 "open_pmd - %s_%s does not seem to have " 220 "unitDimension", 221 pname, 222 recname, 223 ) 224 for i in self.known_particle_fields: 225 mylog.debug("open_pmd - known_particle_fields - %s", i) 226 except (KeyError, TypeError, AttributeError): 227 pass 228 229 super().__init__(ds, field_list) 230 231 def setup_fluid_fields(self): 232 """Defines which derived mesh fields to create. 233 234 If a field can not be calculated, it will simply be skipped. 235 """ 236 # Set up aliases first so the setup for poynting can use them 237 if len(self._mag_fields) > 0: 238 setup_magnetic_field_aliases(self, "openPMD", self._mag_fields) 239 setup_poynting_vector(self) 240 241 def setup_particle_fields(self, ptype): 242 """Defines which derived particle fields to create. 243 244 This will be called for every entry in 245 `OpenPMDDataset``'s ``self.particle_types``. 246 If a field can not be calculated, it will simply be skipped. 247 """ 248 setup_absolute_positions(self, ptype) 249 setup_kinetic_energy(self, ptype) 250 setup_velocity(self, ptype) 251 super().setup_particle_fields(ptype) 252