Adding another gufunc for sgp4_tsince

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);

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)

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}
 };
 

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,
+)

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