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