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

tests/core/test_spacecraft.py

@@ -0,0 +1,210 @@
+#!/usr/bin/env python
+"""Tests for spacecraft handling."""
+import numpy as np
+import pandas as pd
+import pandas.testing as pdt
+from scipy import constants
+from unittest import TestCase
+from abc import ABC, abstractclassmethod, abstractproperty
+
+# import test_base as base
+from . import test_base as base
+
+from solarwindpy import vector
+from solarwindpy import spacecraft
+
+pd.set_option("mode.chained_assignment", "raise")
+
+
+class TestBase(ABC):
+    @classmethod
+    def setUpClass(cls):
+        data = base.TestData()
+        #         pdb.set_trace()
+        cls.data = data.spacecraft_data
+        cls.set_object_testing()
+
+    #         print("TestBase.setUpClass", flush=True)
+    #         test_plasma = {
+    #             ("pos_HCI", "x", ""): {0: -42, 1: -22, 2: -34},
+    #             ("pos_HCI", "y", ""): {0: 23, 1: 31, 2: 11},
+    #             ("pos_HCI", "z", ""): {0: 35, 1: 27, 2: 49},
+    #             ("v_HCI", "x", ""): {0: 9.0, 1: 10.0, 2: 8.0},
+    #             ("v_HCI", "y", ""): {0: -80.0, 1: -70.0, 2: -90.0},
+    #             ("v_HCI", "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},
+    #             ("gse", "x", ""): {0: 230, 1: 235, 2: 240},
+    #             ("gse", "y", ""): {0: 50, 1: 20, 2: 10},
+    #             ("gse", "z", ""): {0: 30, 1: 25, 2: -50},
+    #         }
+    #
+    #         test_data = pd.DataFrame.from_dict(
+    #             test_plasma, orient="columns", dtype=np.float64
+    #         )
+    #         test_data.columns.names = ["M", "C", "S"]
+    #         cls.data = test_data.xs("", axis=1, level="S")
+    #         cls.set_object_testing()
+    #         print("Done with TestBase", flush=True)
+
+    #         super(TestBase, cls).setUpClass()
+    #         # print(cls.data.iloc[:, :7])
+    #         # print(cls.data.columns.values)
+    #         cls.data = cls.spacecraft_data
+    #         del cls.spacecraft_data
+    #         cls.set_object_testing()
+
+    @abstractclassmethod
+    def set_object_testing(cls):
+        pass
+
+    @abstractproperty
+    def name(self):
+        pass
+
+    @abstractproperty
+    def frame(self):
+        pass
+
+    def test_position(self):
+        cols = pd.Index(("x", "y", "z"), name="C")
+        ot = self.object_testing
+        pdt.assert_index_equal(cols, ot.position.data.columns)
+        self.assertIsInstance(ot.position, vector.Vector)
+        self.assertEqual(ot.position, ot.r)
+        self.assertEqual(ot.position, ot.pos)
+        return ot
+
+    def test_velocity(self):
+        cols = pd.Index(("x", "y", "z"), name="C")
+        ot = self.object_testing
+        pdt.assert_index_equal(cols, ot.velocity.data.columns)
+        self.assertIsInstance(ot.velocity, vector.Vector)
+        self.assertEqual(ot.velocity, ot.v)
+        return ot
+
+    def test_data(self):
+        ot = self.object_testing
+        pdt.assert_frame_equal(self.data, ot.data)
+
+    def test_name(self):
+        ot = self.object_testing
+        self.assertEqual(self.name, ot.name)
+
+    def test_frame(self):
+        ot = self.object_testing
+        self.assertEqual(self.frame, ot.frame)
+
+    def test_distance2sun(self):
+        ot = self.object_testing
+
+        frame = self.frame
+        pos = self.data.loc[:, "pos"]
+        if frame == "GSE":
+            # Origin is Earth, so we need to transform x-component to sun-centered.
+            assert pos.columns.equals(pd.Index(("x", "y", "z"), name="C"))
+            au = constants.au  # 1 AU in meters
+            re = 6378.1e3  # Earth radius in meters
+            sign_x = re * 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 = pos.multiply(sign_x, axis=1).add(change_origin, axis=1)
+
+        elif frame == "HCI":
+            # Origin is sun and propagationd distance is just magnitude
+            assert pos.columns.equals(pd.Index(("x", "y", "z"), name="C"))
+            rs = 695.508e6  # Sun radius in meters
+            pos = pos.multiply(rs)
+
+        else:
+            raise NotImplementedError("No test written for frame {}".format(frame))
+
+        dist = pos.pow(2).sum(axis=1).pipe(np.sqrt)
+        dist.name = "distance2sun"
+        pdt.assert_series_equal(dist, ot.distance2sun)
+
+
+class TestWind(TestBase, TestCase):
+    @classmethod
+    def set_object_testing(cls):
+        data = cls.data.xs("gse", axis=1, level="M")
+        data = pd.concat({"pos": data}, axis=1, names=["M"], sort=True).sort_index(
+            axis=1
+        )
+        cls.data = data
+        sc = spacecraft.Spacecraft(data, "wind", "gse")
+
+        cls.object_testing = sc
+
+    @property
+    def frame(self):
+        return "GSE"
+
+    @property
+    def name(self):
+        return "WIND"
+
+    def test_position(self):
+        super(TestWind, self).test_position()
+        pos = self.data.xs("pos", axis=1, level="M")
+        ot = self.object_testing
+        pdt.assert_frame_equal(pos, ot.position.data)
+
+    def test_velocity(self):
+        with self.assertRaises(KeyError):
+            self.object_testing.velocity
+        with self.assertRaises(KeyError):
+            self.object_testing.v
+
+    def test_carrington(self):
+        with self.assertRaises(KeyError):
+            self.object_testing.carrington
+
+
+class TestPSP(TestBase, TestCase):
+    @classmethod
+    def set_object_testing(cls):
+        p = cls.data.xs("pos_HCI", axis=1, level="M")
+        v = cls.data.xs("v_HCI", axis=1, level="M")
+        c = cls.data.xs("Carr", axis=1, level="M")
+
+        data = pd.concat(
+            {"v": v, "pos": p, "carr": c}, axis=1, names=["M"], sort=True
+        ).sort_index(axis=1)
+        sc = spacecraft.Spacecraft(data, "psp", "hci")
+        cls.object_testing = sc
+        cls.data = data
+
+    @property
+    def frame(self):
+        return "HCI"
+
+    @property
+    def name(self):
+        return "PSP"
+
+    def test_position(self):
+        super(TestPSP, self).test_position()
+
+        pos = self.data.xs("pos", axis=1, level="M")
+        ot = self.object_testing
+        pdt.assert_frame_equal(pos, ot.position.data)
+
+    def test_velocity(self):
+        super(TestPSP, self).test_velocity()
+
+        v = self.data.xs("v", axis=1, level="M")
+        ot = self.object_testing
+        pdt.assert_frame_equal(v, ot.velocity.data)
+
+    def test_carrington(self):
+        cols = pd.Index(("lat", "lon"), name="C")
+        carr = self.data.xs("carr", axis=1, level="M")
+
+        ot = self.object_testing
+        self.assertIsInstance(ot.carrington, pd.DataFrame)
+        pdt.assert_index_equal(cols, ot.carrington.columns)
+        pdt.assert_frame_equal(carr, ot.carrington)

tests/core/test_units_constants.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python
+"""Tests for units and constants containers."""
+
+import pandas as pd
+
+from solarwindpy.core import units_constants as uc
+
+
+def test_units_attributes():
+    u = uc.Units()
+    assert hasattr(u, "b")
+    assert isinstance(u.b, float)
+    assert hasattr(u, "v")
+    assert isinstance(u.v, float)
+
+
+def test_constants_attributes():
+    c = uc.Constants()
+    assert hasattr(c, "kb")
+    assert isinstance(c.kb, pd.Series)
+    assert hasattr(c, "misc")
+    assert isinstance(c.misc, pd.Series)

tests/data/epoch.csv

@@ -0,0 +1,4 @@
+epoch
+1995-01-01T12:35:00
+2022-03-23T19:29:09
+2022-10-09T01:47:01.23456

tests/data/plasma.csv

@@ -0,0 +1,4 @@
+b|x|,b|y|,b|z|,b|theta_rms|,b|mag_rms|,n||a,n||p1,n||p2,v|x|a,v|x|p1,v|x|p2,v|y|a,v|y|p1,v|y|p2,v|z|a,v|z|p1,v|z|p2,w|par|a,w|par|p1,w|par|p2,w|per|a,w|per|p1,w|per|p2,n||e,v|x|e,v|y|e,v|z|e,w|par|e,w|per|e
+0.5,-0.25,0.3,0.01,0.1,0.5,1.0,2.0,125.0,100.0,150.0,250.0,200.0,300.0,500.0,400.0,600.0,3.0,10.0,7.5,7.0,7.0,5.5,2.0,130.0,230.0,700.0,5.0,6.0
+0.6,-0.26,0.4,0.07,0.2,1.0,2.0,3.0,250.0,200.0,300.0,375.0,300.0,450.0,750.0,600.0,900.0,4.0,20.0,15.0,9.0,26.0,20.0,3.0,230.0,450.0,800.0,12.0,22.0
+0.7,0.27,-0.7,0.03,0.3,1.5,3.0,4.0,375.0,300.0,450.0,750.0,600.0,900.0,1000.0,800.0,1200.0,5.0,30.0,22.5,10.0,28.0,30.5,7.0,380.0,800.0,900.0,25.0,15.0

tests/data/spacecraft.csv

@@ -0,0 +1,4 @@
+pos_HCI|x|,pos_HCI|y|,pos_HCI|z|,v_HCI|x|,v_HCI|y|,v_HCI|z|,Carr|lat|,Carr|lon|,gse|x|,gse|y|,gse|z|
+-42,23,35,9.0,-80.0,-0.5,-2.0,-26.0,230,50,30
+-22,31,27,10.0,-70.0,0.5,-1.0,-36.0,235,20,25
+-34,11,49,8.0,-90.0,1.5,3.0,-16.0,240,10,-50

tests/fitfunctions/conftest.py

@@ -0,0 +1,60 @@
+"""Shared fixtures for fitfunction tests."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+
+@pytest.fixture
+def simple_linear_data():
+    """Noisy linear data with unit weights.
+
+    Returns
+    -------
+    tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]
+        Tuple ``(x, y, w)`` where ``y`` follows ``2*x + 1`` with added
+        Gaussian noise and ``w`` contains unit weights.
+    """
+    rng = np.random.default_rng(42)
+    x = np.linspace(0.0, 1.0, 20)
+    noise = rng.normal(scale=0.05, size=x.shape)
+    y = 2.0 * x + 1.0 + noise
+    w = np.ones_like(x)
+    return x, y, w
+
+
+@pytest.fixture
+def gauss_data():
+    """Synthetic Gaussian data.
+
+    Returns
+    -------
+    tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]
+        Tuple ``(x, y, w)`` where ``y`` is a Gaussian with additive noise
+        and ``w`` contains unit weights.
+    """
+    rng = np.random.default_rng(42)
+    amplitude = 1.0
+    mean = 0.0
+    sigma = 1.0
+    x = np.linspace(-5.0, 5.0, 100)
+    noise = rng.normal(scale=0.05, size=x.shape)
+    y = amplitude * np.exp(-0.5 * ((x - mean) / sigma) ** 2) + noise
+    w = np.ones_like(x)
+    return x, y, w
+
+
+@pytest.fixture
+def small_n():
+    """Minimal dataset insufficient for fitting.
+
+    Returns
+    -------
+    tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]
+        Tuple ``(x, y, w)`` containing a single observation.
+    """
+    x = np.linspace(0.0, 1.0, 1)
+    y = 2.0 * x + 1.0
+    w = np.ones_like(x)
+    return x, y, w

tests/fitfunctions/test_core.py

@@ -0,0 +1,193 @@
+import numpy as np
+import pytest
+from types import SimpleNamespace
+
+from solarwindpy.fitfunctions.core import (
+    FitFunction,
+    ChisqPerDegreeOfFreedom,
+    InitialGuessInfo,
+)
+
+
+def linear_function(x, m, b):
+    return m * x + b
+
+
+class LinearFit(FitFunction):
+    @property
+    def function(self):
+        return linear_function
+
+    @property
+    def p0(self):
+        return [0.0, 0.0]
+
+    @property
+    def TeX_function(self):
+        return "m x + b"
+
+
+def test_clean_raw_obs():
+    lf = LinearFit([0, 1], [1, 2])
+    with pytest.raises(ValueError):
+        lf._clean_raw_obs([0, 1], [1], None)
+    x, y, w = lf._clean_raw_obs([0, 1], [1, 2], [1, 1])
+    assert np.array_equal(x, np.array([0, 1]))
+    assert np.array_equal(y, np.array([1, 2]))
+    assert np.array_equal(w, np.array([1, 1]))
+
+
+def test_build_one_obs_mask():
+    lf = LinearFit([0, 1], [1, 2])
+    x = np.array([0.0, 1.0, 2.0, np.nan])
+    mask = lf._build_one_obs_mask("xobs", x, 0.5, 1.5)
+    assert np.array_equal(mask, np.array([False, True, False, False]))
+    mask = lf._build_one_obs_mask("xobs", x, None, 1.5)
+    assert np.array_equal(mask, np.array([True, True, False, False]))
+    mask = lf._build_one_obs_mask("xobs", x, 0.5, None)
+    assert np.array_equal(mask, np.array([False, True, True, False]))
+
+
+def test_build_outside_mask():
+    lf = LinearFit([0, 1], [1, 2])
+    x = np.arange(5)
+    mask = lf._build_outside_mask("x", x, None)
+    assert np.array_equal(mask, np.ones_like(x, dtype=bool))
+    mask = lf._build_outside_mask("x", x, (1, 3))
+    assert np.array_equal(mask, np.array([True, True, False, True, True]))
+
+
+def test_set_fit_obs_combined_masks():
+    x = np.array([0, 1, 2, 3, 4], dtype=float)
+    y = np.array([10, 20, 30, 40, 50], dtype=float)
+    w = np.array([0.5, 1, 2, 3, 4], dtype=float)
+    lf = LinearFit(x, y, weights=w)
+    lf.set_fit_obs(
+        x,
+        y,
+        w,
+        xmin=1,
+        xmax=3,
+        ymin=15,
+        ymax=45,
+        wmin=0.02,
+        wmax=0.03,
+        logy=True,
+    )
+    assert np.array_equal(lf.observations.used.x, np.array([1.0, 2.0]))
+    assert np.array_equal(lf.observations.used.y, np.array([20.0, 30.0]))
+    assert np.array_equal(lf.observations.used.w, np.array([1.0, 2.0]))
+    assert np.array_equal(
+        lf.observations.tk_observed, np.array([False, True, True, False, False])
+    )
+
+
+def test_set_argnames():
+    lf = LinearFit([0, 1], [1, 2])
+    assert lf.argnames == ["m", "b"]
+
+
+def test_run_least_squares(monkeypatch, simple_linear_data):
+    x, y, w = simple_linear_data
+    lf = LinearFit(x, y, weights=w)
+    captured = {}
+
+    def fake_ls(func, p0, **kwargs):
+        captured.update(kwargs)
+        jac = np.eye(lf.observations.used.x.size, len(p0))
+        return SimpleNamespace(
+            success=True, x=p0, cost=0.0, jac=jac, fun=np.zeros(lf.nobs)
+        )
+
+    from solarwindpy.fitfunctions import core as core_module
+
+    monkeypatch.setattr(core_module, "least_squares", fake_ls)
+    res, p0 = lf._run_least_squares()
+    assert captured["method"] == "trf"
+    assert captured["loss"] == "huber"
+    assert captured["max_nfev"] == 10000
+    assert captured["f_scale"] == 0.1
+    assert captured["jac"] == "2-point"
+    assert np.array_equal(p0, np.array(lf.p0))
+
+    with pytest.raises(ValueError):
+        lf._run_least_squares(args=(1,))
+
+
+def test_calc_popt_pcov_psigma_chisq():
+    x = np.array([0.0, 1.0, 2.0])
+    y = 2.0 * x + 1.0
+    lf = LinearFit(x, y)
+    jac = np.array([[0.0, 1.0], [1.0, 1.0], [2.0, 1.0]])
+    res = SimpleNamespace(
+        x=np.array([2.0, 1.0]), cost=0.0, fun=np.zeros_like(x), jac=jac
+    )
+    popt, pcov, psigma, chisq = lf._calc_popt_pcov_psigma_chisq(
+        res, p0=np.array([0.0, 0.0])
+    )
+    assert np.allclose(popt, np.array([2.0, 1.0]))
+    assert np.array_equal(pcov, np.zeros((2, 2)))
+    assert np.array_equal(psigma, np.zeros(2))
+    assert chisq.linear == 0.0 and chisq.robust == 0.0
+
+
+def test_make_fit_success_failure(monkeypatch, simple_linear_data, small_n):
+    x, y, w = simple_linear_data
+    lf = LinearFit(x, y, weights=w)
+    lf.make_fit()
+    assert isinstance(lf.fit_result, object)
+    assert set(lf.popt) == {"m", "b"}
+    assert set(lf.psigma) == {"m", "b"}
+    assert lf.pcov.shape == (2, 2)
+    assert isinstance(lf.chisq_dof, ChisqPerDegreeOfFreedom)
+    assert lf.plotter is not None and lf.TeX_info is not None
+
+    x, y, w = small_n
+    lf_small = LinearFit(x, y, weights=w)
+    err = lf_small.make_fit(return_exception=True)
+    assert isinstance(err, ValueError)
+
+    def fail_run(*_, **__):
+        raise RuntimeError("fail")
+
+    x, y, w = simple_linear_data
+    lf_fail = LinearFit(x, y, weights=w)
+    monkeypatch.setattr(LinearFit, "_run_least_squares", fail_run)
+    err = lf_fail.make_fit(return_exception=True)
+    assert isinstance(err, RuntimeError)
+    with pytest.raises(RuntimeError):
+        lf_fail.make_fit()
+
+
+@pytest.fixture
+def fitted_linear(simple_linear_data):
+    x, y, w = simple_linear_data
+    lf = LinearFit(x, y, weights=w)
+    lf.make_fit()
+    return lf
+
+
+def test_str_call_and_properties(fitted_linear):
+    lf = fitted_linear
+    s = str(lf)
+    assert "LinearFit" in s and "m x + b" in s
+    xnew = np.array([0.0, 0.5])
+    ypred = lf(xnew)
+    assert np.allclose(ypred, lf.popt["m"] * xnew + lf.popt["b"], rtol=1e-2, atol=1e-2)
+    assert lf.argnames == ["m", "b"]
+    assert isinstance(lf.fit_bounds, dict)
+    assert isinstance(lf.chisq_dof, ChisqPerDegreeOfFreedom)
+    assert lf.dof == lf.observations.used.y.size - len(lf.p0)
+    assert lf.fit_result is not None
+    assert isinstance(lf.initial_guess_info["m"], InitialGuessInfo)
+    assert lf.nobs == lf.observations.used.x.size
+    assert lf.plotter is not None
+    assert set(lf.popt) == {"m", "b"}
+    assert set(lf.psigma) == {"m", "b"}
+    assert set(lf.psigma_relative) == {"m", "b"}
+    combined = lf.combined_popt_psigma
+    assert set(combined) == {"popt", "psigma", "psigma_relative"}
+    assert lf.pcov.shape == (2, 2)
+    assert 0.0 <= lf.rsq <= 1.0
+    assert lf.sufficient_data is True
+    assert lf.TeX_info is not None