From 5f6e9881db2c41d1877abe17166543a16e1a5ea2 Mon Sep 17 00:00:00 2001 From: "Benjamin V. Root" Date: Thu, 31 Jul 2025 14:15:52 -0400 Subject: [PATCH] Adding another gufunc for the sgp4_tsince --- extension/_sgp4.cpp | 119 ++++++++++++++++++++++++++++++++++++- extension/wrapper.cpp | 2 + extension/wrapper.h | 6 +- lib/sgp4_ufunc/__init__.py | 4 +- tests/test_sgp4.py | 28 ++++++++- 5 files changed, 150 insertions(+), 9 deletions(-) diff --git a/extension/_sgp4.cpp b/extension/_sgp4.cpp index 61bfb76..e787bcf 100644 --- a/extension/_sgp4.cpp +++ b/extension/_sgp4.cpp @@ -30,6 +30,22 @@ " Last dimension is size 3.\n" \ "" +#define SGP4_TSINCE_GUFUNC_DOCSTRING \ +"Compute positions and velocities for the times in a NumPy array.\n" \ +"\n" \ +"Inputs:\n" \ +" ``satrec`` - Satrec instance(s)\n" \ +" ``tsince`` - (double) Minutes since epoch time in satrec\n" \ +"\n" \ +"Returns:\n" \ +" ``e``: nonzero for any dates that produced errors, 0 otherwise.\n" \ +" ``r``: NumPy array of position vector(s) in kilometers.\n" \ +" Last dimension is size 3.\n" \ +" ``v``: NumPy array of velocity vector(s) in kilometers per second.\n" \ +" Last dimension is size 3.\n" \ +"" + + static inline void set(double *px, ptrdiff_t stride, ptrdiff_t index, double value) { @@ -50,7 +66,6 @@ static void sgp4_core( double r[3], v[3]; double tsince = (*p_jd - satrec.jdsatepoch) * 1440.0 + (*p_fr - satrec.jdsatepochF) * 1440.0; - std::cout << "satnum: " << p_satrec->satrec.satnum[0] << " satepoch: " << satrec.jdsatepoch << " frac: " << satrec.jdsatepochF << " tsince: " << tsince << std::endl; SGP4Funcs::sgp4(satrec, tsince, r, v); if (satrec.error && satrec.error < 6) { r[0] = r[1] = r[2] = v[0] = v[1] = v[2] = NAN; @@ -64,6 +79,31 @@ static void sgp4_core( set(p_vel, vel_stride, 2, v[2]); } +static void sgp4_tsince_core( + SatrecObject *p_satrec, + double *p_tsince, + uint8_t *p_err, + double *p_pos, + npy_intp pos_stride, + double *p_vel, + npy_intp vel_stride +) { + elsetrec satrec = p_satrec->satrec; + double r[3], v[3]; + SGP4Funcs::sgp4(satrec, *p_tsince, r, v); + if (satrec.error && satrec.error < 6) { + r[0] = r[1] = r[2] = v[0] = v[1] = v[2] = NAN; + } + p_err[0] = satrec.error; + set(p_pos, pos_stride, 0, r[0]); + set(p_pos, pos_stride, 1, r[1]); + set(p_pos, pos_stride, 2, r[2]); + set(p_vel, vel_stride, 0, v[0]); + set(p_vel, vel_stride, 1, v[1]); + set(p_vel, vel_stride, 2, v[2]); +} + + // // Instantiated loop functions with C calling convention // @@ -127,6 +167,64 @@ static PyUFuncGenericFunction funcs[] = { &loop_sgp4_core }; + +extern "C" { + +static void loop_sgp4_tsince_core( + char **args, const npy_intp *dimensions, + const npy_intp* steps, void* data) +{ + char *p_satrec = args[0]; + char *p_tsince = args[1]; + char *p_out1 = args[2]; + char *p_out2 = args[3]; + char *p_out3 = args[4]; + npy_intp nloops = dimensions[0]; + + assert(dimensions[1] == 3); + for (int j = 0; j < nloops; + ++j, + p_satrec += steps[0], + p_tsince += steps[1], + p_out1 += steps[2], + p_out2 += steps[3], + p_out3 += steps[4]) { + PyObject** p_item = (PyObject **) p_satrec; + if (!PyObject_IsInstance(*p_item, (PyObject*) &SatrecType)) { + NPY_ALLOW_C_API_DEF + NPY_ALLOW_C_API + PyErr_Format(PyExc_TypeError, "satrec must be a Satrec or an array of Satrec," + " but element %d is: %R", j, *p_item); + NPY_DISABLE_C_API + return; + } + + sgp4_tsince_core( + (SatrecObject *) *p_item, + (npy_double *) p_tsince, + (npy_ubyte *) p_out1, + (npy_double *) p_out2, steps[5], + (npy_double *) p_out3, steps[6]); + } +} +} // extern "C" + +/* These are the input and return dtypes of sgp4_tsince.*/ +static const char sgp4_tsince_types[] = { + // Inputs + // TODO: Maybe make this object be a custom dtype? + NPY_OBJECT, NPY_DOUBLE, + // Outputs + NPY_UINT8, NPY_DOUBLE, NPY_DOUBLE +}; + + +/*This a pointer to the above *loop* function*/ +static PyUFuncGenericFunction tsince_funcs[] = { + &loop_sgp4_tsince_core +}; + + #define MODULE_DOCSTRING \ "SGP4\n" \ "\n" \ @@ -152,7 +250,7 @@ static struct PyModuleDef moduledef = { PyMODINIT_FUNC PyInit__sgp4(void) { - PyObject *m, *_sgp4_gufunc, *d; + PyObject *m, *_sgp4_gufunc, *_sgp4_tsince_gufunc, *d; SatrecType.tp_name = "Satrec"; SatrecType.tp_basicsize = sizeof(SatrecObject); @@ -181,13 +279,30 @@ PyMODINIT_FUNC PyInit__sgp4(void) "(),(),()->(),(3),(3)"); if (_sgp4_gufunc == NULL) { + Py_DECREF(m); + return NULL; + } + + _sgp4_tsince_gufunc = PyUFunc_FromFuncAndDataAndSignature( + tsince_funcs, NULL, sgp4_tsince_types, 1, 2, 3, + PyUFunc_None, "_sgp4_tsince_gufunc", + SGP4_TSINCE_GUFUNC_DOCSTRING, 0, + "(),()->(),(3),(3)"); + + if (_sgp4_tsince_gufunc == NULL) { + Py_DECREF(_sgp4_gufunc); + Py_DECREF(m); return NULL; } + d = PyModule_GetDict(m); PyDict_SetItemString(d, "_sgp4_gufunc", _sgp4_gufunc); Py_DECREF(_sgp4_gufunc); + PyDict_SetItemString(d, "_sgp4_tsince_gufunc", _sgp4_tsince_gufunc); + Py_DECREF(_sgp4_tsince_gufunc); + Py_INCREF(&SatrecType); if (PyModule_AddObject(m, "Satrec", (PyObject *) &SatrecType) < 0) { Py_DECREF(&SatrecType); diff --git a/extension/wrapper.cpp b/extension/wrapper.cpp index 4074ee1..bb69e14 100644 --- a/extension/wrapper.cpp +++ b/extension/wrapper.cpp @@ -159,6 +159,7 @@ Satrec_sgp4init(PyObject *self, PyObject *args) return Py_None; } +/* PyObject * Satrec_sgp4_tsince(PyObject *self, PyObject *args) { @@ -173,6 +174,7 @@ Satrec_sgp4_tsince(PyObject *self, PyObject *args) return Py_BuildValue("i(fff)(fff)", satrec.error, r[0], r[1], r[2], v[0], v[1], v[2]); } +*/ PyObject * Satrec_sgp4(PyObject *self, PyObject *args) diff --git a/extension/wrapper.h b/extension/wrapper.h index afbc53b..1b385df 100644 --- a/extension/wrapper.h +++ b/extension/wrapper.h @@ -18,7 +18,7 @@ static PyTypeObject SatrecType = { PyObject *Satrec_twoline2rv(PyTypeObject *cls, PyObject *args); PyObject *Satrec_sgp4init(PyObject *self, PyObject *args); PyObject *Satrec_sgp4(PyObject *self, PyObject *args); -PyObject *Satrec_sgp4_tsince(PyObject *self, PyObject *args); +//PyObject *Satrec_sgp4_tsince(PyObject *self, PyObject *args); static PyMethodDef Satrec_methods[] = { {"twoline2rv", (PyCFunction)Satrec_twoline2rv, METH_VARARGS | METH_CLASS, @@ -29,8 +29,8 @@ static PyMethodDef Satrec_methods[] = { PyDoc_STR("Given a modified julian date, return position and velocity.")}, //{"_sgp4", (PyCFunction)Satrec__sgp4, METH_VARARGS, // PyDoc_STR("Given an array of modified julian dates, return position and velocity arrays.")}, - {"sgp4_tsince", (PyCFunction)Satrec_sgp4_tsince, METH_VARARGS, - PyDoc_STR("Given minutes since epoch, return position and velocity.")}, + //{"sgp4_tsince", (PyCFunction)Satrec_sgp4_tsince, METH_VARARGS, + // PyDoc_STR("Given minutes since epoch, return position and velocity.")}, {NULL, NULL} }; diff --git a/lib/sgp4_ufunc/__init__.py b/lib/sgp4_ufunc/__init__.py index 1a2df5f..a992e6c 100644 --- a/lib/sgp4_ufunc/__init__.py +++ b/lib/sgp4_ufunc/__init__.py @@ -1 +1,3 @@ -from ._sgp4 import _sgp4_gufunc, Satrec, WGS72, WGS72OLD, WGS84 +from ._sgp4 import ( + _sgp4_gufunc, _sgp4_tsince_gufunc, Satrec, WGS72, WGS72OLD, WGS84, +) diff --git a/tests/test_sgp4.py b/tests/test_sgp4.py index 888f8bb..823aaa1 100644 --- a/tests/test_sgp4.py +++ b/tests/test_sgp4.py @@ -4,7 +4,7 @@ import pytest -from sgp4_ufunc import _sgp4_gufunc, Satrec, WGS72 +from sgp4_ufunc import _sgp4_gufunc, _sgp4_tsince_gufunc, Satrec, WGS72 def sgp4init_args(d): """Given a dict of orbital parameters, return them in sgp4init order.""" @@ -74,6 +74,18 @@ def test_smoketest(basic_satellite): assert v.dtype == np.dtype('float64') assert v.shape == (10, 4, 3) + +def test_smoketest_tsince(basic_satellite): + tsince = np.zeros((1, 4)) + e, r, v, = _sgp4_tsince_gufunc(basic_satellite, tsince) + assert e.dtype == np.dtype('uint8') + assert e.shape == (1, 4) + assert r.dtype == np.dtype('float64') + assert r.shape == (1, 4, 3) + assert v.dtype == np.dtype('float64') + assert v.shape == (1, 4, 3) + + @pytest.mark.parametrize("sat,jd,fr", [ (Satrec.twoline2rv(LINE1, LINE2), np.zeros((4,)), np.zeros((4,), dtype='S')), (Satrec.twoline2rv(LINE1, LINE2), ['foobar'], np.zeros((4,))), @@ -85,6 +97,18 @@ def test_invalid_inputs(sat, jd, fr): _sgp4_gufunc(sat, jd, fr) +@pytest.mark.parametrize("sat,tsince", [ + (Satrec.twoline2rv(LINE1, LINE2), np.zeros((4,), dtype='S')), + (Satrec.twoline2rv(LINE1, LINE2), ['foobar']), + (None, np.zeros((4,))), + ([Satrec.twoline2rv(LINE1, LINE2), datetime(2025, 7, 1)], np.zeros((2,))), +]) +def test_invalid_inputs_tsince(sat, tsince): + + with pytest.raises(TypeError): + _sgp4_tsince_gufunc(sat, tsince) + + def test_whether_array_logic_writes_nan_values_to_correct_row(): # https://github.com/brandon-rhodes/python-sgp4/issues/87 l1 = "1 44160U 19006AX 20162.79712247 +.00816806 +19088-3 +34711-2 0 9997" @@ -97,8 +121,6 @@ def test_whether_array_logic_writes_nan_values_to_correct_row(): assert np.isnan(r).tolist() == [[False, False, False], [True, True, True]] assert np.isnan(v).tolist() == [[False, False, False], [True, True, True]] - - def verify_vanguard_1(sat, legacy=False): attrs = VANGUARD_ATTRS