solarwindpy 0.0.1.dev0__py3-none-any.whl → 0.1.1__py3-none-any.whl

solarwindpy/core/spacecraft.py

@@ -0,0 +1,256 @@
+#!/usr/bin/env python
+r"""Contais :py:class:`~solarwindpy.core.spacecraft.Spacecraft` class.
+
+Class inherets from :py:class:`~solarwindpy.core.base.Base` and contains :py:class:`~solarwindpy.core.vector.Vector` objects.
+"""
+
+import pandas as pd
+import numpy as np
+
+# We rely on views via DataFrame.xs to reduce memory size and do not
+# `.copy(deep=True)`, so we want to make sure that this doesn't
+# accidentally cause a problem.
+
+from . import base
+from . import vector
+
+
+class Spacecraft(base.Base):
+    r"""Representation of a spacecraft trajectory.
+
+    Parameters
+    ----------
+    data : :class:`pandas.DataFrame`
+        Vector position (and optionally velocity) with MultiIndex columns
+        ``("M", "C")``.
+    name : str
+        Identifier of the spacecraft.
+    frame : str
+        Reference frame of the vectors, e.g. ``"HCI"`` or ``"GSE"``.
+    """
+
+    def __init__(self, data, name, frame):
+        r"""Initialize a spacecraft with `data`.
+
+        Parameters
+        ----------
+        data: pd.DataFrame
+            2-level MultiIndex columns identifying measurment "M" and component "C".
+            Should contain at a minimum vector position. Can also contain vector
+            velocity and Carrington location. If vector velocity included, it should be
+            in the same frame of reference as the position.
+        name: str
+            Identify the spacecraft, e.g. Parker Solar Probe (PSP) or Wind. Internally
+            stored in all caps for consistency.
+        frame: str
+            The frame of reference for the spacecraft position and velocity, e.g.
+            Geocentric Solar Ecliptic (GSE) or Heliocentric Internal (HCI).
+
+        Examples
+        --------
+        >>> epoch = pd.Series({0: pd.to_datetime("1995-01-01"),
+                               1: pd.to_datetime("2015-03-23"),
+                               2: pd.to_datetime("2022-10-09")}, name="Epoch")
+        >>> data = {("pos", "x", ""): {0: -42, 1: -22, 2: -34},
+                    ("pos", "y", ""): {0: 23, 1: 31, 2: 11},
+                    ("pos", "z", ""): {0: 35, 1: 27, 2: 49},
+                    ("v", "x", ""): {0: 9.0, 1: 10.0, 2: 8.0},
+                    ("v", "y", ""): {0: -80.0, 1: -70.0, 2: -90.0},
+                    ("v", "z", ""): {0: -0.5, 1: 0.5, 2: 1.5},
+                    ("carr", "lat", ""): {0: -2.0, 1: -1.0, 2: 3.0},
+                    ("carr", "lon", ""): {0: -26.0, 1: -36.0, 2: -16.0}}
+        >>> spacecraft = pd.DataFrame.from_dict(data,
+                                                orient="columns",
+                                                dtype=np.float64)
+        >>> spacecraft.index = epoch
+        >>> spacecraft.columns.names = ["M", "C", "S"]
+        >>> spacecraft = spacecraft.xs("", axis=1, level="S")
+        >>> spacecraft
+        M            pos                 v            carr
+        C              x     y     z     x     y    z  lat   lon
+        Epoch
+        1995-01-01 -42.0  23.0  35.0   9.0 -80.0 -0.5 -2.0 -26.0
+        2015-03-23 -22.0  31.0  27.0  10.0 -70.0  0.5 -1.0 -36.0
+        2022-10-09 -34.0  11.0  49.0   8.0 -90.0  1.5  3.0 -16.0
+        >>> spacecraft = Spacecraft(spacecraft, "PSP", "HCI")
+        """
+        super(Spacecraft, self).__init__(data)
+        self.set_frame_name(frame, name)
+        self.set_data(data)
+        self._log_spacecraft()
+
+    @property
+    def frame(self):
+        r"""Spacecraft's frame of reference (e.g. GSE, HCI, etc.)."""
+        return self._frame
+
+    @property
+    def name(self):
+        r"""Spacecraft name (e.g. WIND, PSP)."""
+        return self._name
+
+    @property
+    def position(self):
+        """Position vector of the spacecraft.
+
+        Returns
+        -------
+        vector.Vector
+            Position vector with x, y, z components.
+        """
+        pos = self.data.xs("pos", axis=1, level="M").loc[:, ("x", "y", "z")]
+        return vector.Vector(pos)
+
+    @property
+    def pos(self):
+        """Shortcut to position property.
+
+        Returns
+        -------
+        vector.Vector
+            Position vector with x, y, z components.
+        """
+        # Ensures that `sc.pos` returns vector.
+        return self.position
+
+    @property
+    def r(self):
+        r"""Shortcut to :py:attr:`position`."""
+        return self.position
+
+    @property
+    def velocity(self):
+        """Velocity vector of the spacecraft.
+
+        Returns
+        -------
+        vector.Vector
+            Velocity vector with x, y, z components.
+
+        Raises
+        ------
+        KeyError
+            If spacecraft velocity data is not available.
+        """
+        try:
+            v = self.data.xs("v", axis=1, level="M").loc[:, ("x", "y", "z")]
+            return vector.Vector(v)
+        except KeyError as e:  # noqa: F841
+            raise KeyError("Spacecraft doesn't know it's velocity.")
+
+    @property
+    def v(self):
+        r"""Shortcut to :py:attr:`velocity`."""
+        return self.velocity
+
+    @property
+    def carrington(self):
+        r"""Carrington latitude and longitude."""
+        try:
+            return self.data.xs("carr", axis=1, level="M").loc[:, ("lat", "lon")]
+        except KeyError as e:  # noqa: F841
+            raise KeyError("Spacecraft doesn't know its Carrington location.")
+
+    @property
+    def distance2sun(self):
+        r"""Radial distance to Sun in meters."""
+        pos = self.position.data
+        frame = self.frame
+
+        if frame == "GSE":
+            re = self.constants.misc.loc["Re [m]"]
+            au = self.constants.misc.loc["1AU [m]"]
+            sign_x = pd.Series(
+                [-1.0, 1.0, 1.0], index=pd.Index(("x", "y", "z"), name="C")
+            )
+            change_origin = pd.Series(
+                [au, 0.0, 0.0], index=pd.Index(("x", "y", "z"), name="C")
+            )
+            pos_SI = (
+                pos.multiply(sign_x, axis=1).multiply(re).add(change_origin, axis=1)
+            )
+
+        elif frame == "HCI":
+            rs = self.constants.misc.loc["Rs [m]"]
+            pos_SI = pos.multiply(rs)
+
+        else:
+            raise NotImplementedError("Unrecognized reference frame `{}`".format(frame))
+
+        # distance2sun units should be [m], so this shouldn't matter. However, just as
+        # beta is treated in this way, we similarly treat distance2sun.
+        d2s = pos_SI.pow(2).sum(axis=1).pipe(np.sqrt) / self.units.distance2sun
+        d2s.name = "distance2sun"
+        return d2s
+
+    def _log_spacecraft(self):
+        self.logger.info(
+            "Created %s spacecraft with %s reference frame", self.name, self.frame
+        )
+
+    def set_frame_name(self, frame, name):
+        """Set the coordinate frame and spacecraft name.
+
+        Parameters
+        ----------
+        frame : str
+            Coordinate frame ('GSE' or 'HCI').
+        name : str
+            Spacecraft name.
+
+        Raises
+        ------
+        NotImplementedError
+            If frame is not 'GSE' or 'HCI'.
+        """
+        frame = frame.upper()
+        name = name.upper()
+
+        if frame not in ("GSE", "HCI"):
+            raise NotImplementedError("Unrecognized frame: {}".format(frame))
+        if name not in ("WIND", "PSP"):
+            raise NotImplementedError("Unrecognized name: {}".format(name))
+
+        self._frame = frame
+        self._name = name
+
+    def set_data(self, data):
+        """Set the spacecraft data.
+
+        Parameters
+        ----------
+        data : pd.DataFrame
+            Spacecraft position/velocity data.
+        """
+        super(Spacecraft, self).set_data(data)
+
+        p = data.xs("pos", axis=1, level="M")
+        #         assert isinstance(p, pd.DataFrame)
+
+        p = p.loc[:, ["x", "y", "z"]]
+        assert p.shape[1] == 3
+
+        target = {"pos": p}
+
+        try:
+            v = data.xs("v", axis=1, level="M")
+            assert isinstance(v, pd.DataFrame)
+            v = v.loc[:, ["x", "y", "z"]]
+            assert v.shape[1] == 3
+            target["v"] = v
+        except KeyError:
+            pass
+
+        try:
+            c = data.xs("carr", axis=1, level="M")
+            assert isinstance(c, pd.DataFrame)
+            c = c.loc[:, ["lat", "lon"]]
+            assert c.shape[1] == 2
+            target["carr"] = c
+        except KeyError:
+            pass
+
+        target = pd.concat(target, axis=1, names=["M"], sort=False).sort_index(axis=1)
+
+        assert isinstance(target.index, pd.DatetimeIndex)
+        self._data = target

solarwindpy/core/tensor.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python
+"""Tensor class for storing quantities like thermal speed, pressure, and temperature."""
+
+import pandas as pd
+
+from . import base
+
+
+class Tensor(base.Base):
+    """Container for tensor-valued quantities.
+
+    Parameters
+    ----------
+    data : :class:`pandas.DataFrame`
+        Tensor data with components ``par``, ``per`` and ``scalar``.
+    """
+
+    def __init__(self, data: pd.DataFrame):
+        """Initialize the Tensor object.
+
+        Parameters
+        ----------
+        data : :class:`pandas.DataFrame`
+            Tensor data to be stored.
+        """
+        super().__init__(data)
+        self._validate_data(data)
+        self._data = data
+
+    def __call__(self, component: str) -> pd.Series | pd.DataFrame:
+        """Access a specific component of the tensor.
+
+        Parameters
+        ----------
+        component : str
+            The name of the component to access.
+
+        Returns
+        -------
+        pd.Series | pd.DataFrame
+            The requested component data.
+
+        Raises
+        ------
+        AttributeError
+            If the component does not exist.
+        """
+        return self.__getattr__(component)
+
+    def set_data(self, new: pd.DataFrame):
+        """Set new tensor data.
+
+        Parameters
+        ----------
+        new : :class:`pandas.DataFrame`
+            The new tensor data.
+
+        Raises
+        ------
+        ValueError
+            If ``new`` does not contain the required columns.
+        """
+        super().set_data(new)
+        self._validate_data(new)
+
+    @staticmethod
+    def _validate_data(data: pd.DataFrame):
+        """Validate the tensor data structure.
+
+        Parameters
+        ----------
+        data : :class:`pandas.DataFrame`
+            Tensor data to validate.
+
+        Raises
+        ------
+        ValueError
+            If the data does not contain the required columns.
+        """
+        required_columns = pd.Index(["per", "par", "scalar"])
+        if not required_columns.isin(data.columns).all():
+            missing_columns = required_columns[~required_columns.isin(data.columns)]
+            raise ValueError(f"Missing required columns: {missing_columns.tolist()}")
+
+    @property
+    def magnitude(self) -> pd.Series | pd.DataFrame:
+        """Calculate and return the magnitude of the tensor."""
+        return self.data.multiply({"par": 1 / 3, "per": 2 / 3}, axis=1, level="C").sum(
+            axis=1
+        )

solarwindpy/core/units_constants.py

@@ -0,0 +1,199 @@
+#!/usr/bin/env python
+"""Definitions of common units and physical constants.
+
+The values are derived from :mod:`scipy.constants`. All quantities stored in
+the :class:`~solarwindpy.core.plasma.Plasma` object have a corresponding entry
+in :class:`Constants` and can be converted using :class:`Units`.
+"""
+
+from dataclasses import dataclass, field
+
+import pandas as pd
+
+from scipy import constants
+from scipy.constants import physical_constants
+
+# We rely on views via DataFrame.xs to reduce memory size and do not
+# `.copy(deep=True)`, so we want to make sure that this doesn't
+# accidentally cause a problem.
+
+_misc_constants = {
+    "e0": constants.epsilon_0,
+    "mu0": constants.mu_0,
+    "c": constants.c,
+    #     "gamma": 5.0 / 3.0,
+    "hbar": physical_constants["Planck constant over 2 pi"][0],
+    "1AU [m]": constants.au,
+    "Re [m]": 6378.1e3,  # Earth Radius in meters
+    "Rs [m]": 695.508e6,  # Sun Radius in meters
+    "gas constant": constants.R,
+}
+
+_kBoltzmann = {
+    "J": constants.k,
+    "eV": physical_constants["Boltzmann constant in eV/K"][0],
+}
+
+_polytropic_index = dict(par=3.0, per=2.0, scalar=5.0 / 3.0)
+
+_m_in_mp = {
+    "p": 1.0,
+    "p1": 1.0,
+    "p2": 1.0,
+    "pm": 1.0,
+    "p_bimax": 1.0,
+    "a": physical_constants["alpha particle-proton mass ratio"][0],
+    "a1": physical_constants["alpha particle-proton mass ratio"][0],
+    "a2": physical_constants["alpha particle-proton mass ratio"][0],
+    "a_bimax": physical_constants["alpha particle-proton mass ratio"][0],
+    "e": physical_constants["electron-proton mass ratio"][0],
+}
+
+_charges = {
+    "e": -constants.e,
+    "p": constants.e,
+    "p1": constants.e,
+    "p2": constants.e,
+    "pm": constants.e,
+    "p_bimax": constants.e,
+    "a": 2.0 * constants.e,
+    "a1": 2.0 * constants.e,
+    "a2": 2.0 * constants.e,
+    "a_bimax": 2.0 * constants.e,
+}
+
+_charge_states = {
+    "e": -1.0,
+    "p": 1.0,
+    "p1": 1.0,
+    "p2": 1.0,
+    "pm": 1.0,
+    "p_bimax": 1.0,
+    "a": 2.0,
+    "a1": 2.0,
+    "a2": 2.0,
+    "a_bimax": 2.0,
+}
+
+_masses = {
+    "e": constants.m_e,
+    "p": constants.m_p,
+    "p1": constants.m_p,
+    "p2": constants.m_p,
+    "pm": constants.m_p,  # proton moment
+    "p_bimax": constants.m_p,
+    "a_bimax": physical_constants["alpha particle mass"][0],
+    "a": physical_constants["alpha particle mass"][0],
+    "a1": physical_constants["alpha particle mass"][0],
+    "a2": physical_constants["alpha particle mass"][0],
+}
+
+_m_amu = {
+    "a": physical_constants["alpha particle mass in u"][0],
+    "a1": physical_constants["alpha particle mass in u"][0],
+    "a2": physical_constants["alpha particle mass in u"][0],
+    "a_bimax": physical_constants["alpha particle mass in u"][0],
+    "p": physical_constants["proton mass in u"][0],
+    "p1": physical_constants["proton mass in u"][0],
+    "p2": physical_constants["proton mass in u"][0],
+    "pm": physical_constants["proton mass in u"][0],
+    "p_bimax": physical_constants["proton mass in u"][0],
+    "e": physical_constants["electron mass in u"][0],
+}
+
+
+@dataclass
+class Constants:
+    """Physical constants useful for solar wind calculations."""
+
+    misc: pd.Series = field(default_factory=lambda: pd.Series(_misc_constants))
+    kb: pd.Series = field(default_factory=lambda: pd.Series(_kBoltzmann))
+    m_in_mp: pd.Series = field(default_factory=lambda: pd.Series(_m_in_mp))
+    m: pd.Series = field(default_factory=lambda: pd.Series(_masses))
+    m_amu: pd.Series = field(
+        default_factory=lambda: pd.Series(_m_amu),
+        metadata={"doc": "Masses in amu."},
+    )
+    charges: pd.Series = field(default_factory=lambda: pd.Series(_charges))
+    charge_states: pd.Series = field(default_factory=lambda: pd.Series(_charge_states))
+    polytropic_index: pd.Series = field(
+        default_factory=lambda: pd.Series(_polytropic_index),
+        metadata={"doc": "Polytropic index for various cases."},
+    )
+
+    def __post_init__(self) -> None:
+        """Validate the shapes of series."""
+        for s in (
+            self.misc,
+            self.kb,
+            self.m_in_mp,
+            self.m,
+            self.m_amu,
+            self.charges,
+            self.charge_states,
+            self.polytropic_index,
+        ):
+            if not isinstance(s, pd.Series):
+                raise TypeError("Constant values must be pandas Series")
+
+
+@dataclass
+class Units:
+    r"""Common unit conversion factors.
+
+    Attributes
+    ----------
+    bfield : float
+        Magnetic field units :math:`[\mathrm{nT}]`.
+    v : float
+        Velocity units :math:`[\mathrm{km\,s^{-1}}]`.
+    pth : float
+        Thermal pressure units :math:`[\mathrm{pPa}]`.
+    temperature : float
+        Temperature units :math:`[10^{5}\,\mathrm{K}]`.
+    n : float
+        Number density units :math:`[\mathrm{cm^{-3}}]`.
+    beta : float
+        Dimensionless beta units.
+    lnlambda : float
+        Dimensionless Coulomb log units.
+    nuc : float
+        Collision frequency units :math:`[10^{-7}\,\mathrm{Hz}]`.
+    nc : float
+        Dimensionless count units.
+    qpar : float
+        Parallel heat flux units :math:`[\mathrm{mW\,cm^{-2}}]`.
+    distance2sun : float
+        Distance to sun units ``[m]``.
+    """
+
+    bfield: float = 1e-9
+    b: float = field(init=False)
+    v: float = 1e3
+    w: float = field(init=False)
+    dv: float = field(init=False)
+    ca: float = field(init=False)
+    cs: float = field(init=False)
+    cfms: float = field(init=False)
+    pth: float = 1e-12
+    temperature: float = 1e5
+    n: float = 1e6
+    rho: float = field(init=False)
+    beta: float = 1.0
+    lnlambda: float = 1.0
+    nuc: float = 1e-7
+    nc: float = 1.0
+    qpar: float = 1e-7
+    distance2sun: float = 1.0
+    specific_entropy: float = field(init=False)
+
+    def __post_init__(self) -> None:
+        """Compute derived unit conversions."""
+        self.b = self.bfield
+        self.w = self.v
+        self.dv = self.v
+        self.ca = self.v
+        self.cs = self.v
+        self.cfms = self.v
+        self.rho = self.n * constants.m_p
+        self.specific_entropy = 1e4 / constants.e