Merge pull request #9 from Rex2000viola/master

Batched mode propagation for satellites with _isimp flag activated

dsgp4/sgp4_batched.py

@@ -102,6 +102,9 @@ def sgp4_batched(satellite, tsince):
     satellite_batch._a=torch.stack([s._a for s in satellite])
     satellite_batch._alta=torch.stack([s._altp for s in satellite])
 
+
+
+
     mrt = torch.zeros(batch_size)
     x2o3  = torch.tensor(2.0 / 3.0)
 
@@ -114,18 +117,15 @@ def sgp4_batched(satellite, tsince):
     xmdf    = satellite_batch._mo + satellite_batch._mdot * satellite_batch._t
     argpdf  = satellite_batch._argpo + satellite_batch._argpdot * satellite_batch._t
     nodedf  = satellite_batch._nodeo + satellite_batch._nodedot * satellite_batch._t
-    argpm   = argpdf
-    mm     = xmdf
+    argpm1   = argpdf
+    mm1     = xmdf
     t2     = satellite_batch._t * satellite_batch._t
     nodem   = nodedf + satellite_batch._nodecf * t2
-    tempa   = 1.0 - satellite_batch._cc1 * satellite_batch._t
-    tempe   = satellite_batch._bstar * satellite_batch._cc4 * satellite_batch._t
-    templ   = satellite_batch._t2cof * t2
+    tempa1   = 1.0 - satellite_batch._cc1 * satellite_batch._t
+    tempe1   = satellite_batch._bstar * satellite_batch._cc4 * satellite_batch._t
+    templ1   = satellite_batch._t2cof * t2
 
-    # START: ASSUME satellite._isimp IS ALWAYS 0
-    for sat in satellite:
-        if sat._isimp == 1:
-            raise ValueError("isimp == 1 not supported in batch mode.")
+
 
     delomg = satellite_batch._omgcof * satellite_batch._t
 
@@ -134,17 +134,22 @@ def sgp4_batched(satellite, tsince):
                 (delmtemp * delmtemp * delmtemp -
                 satellite_batch._delmo)
     temp   = delomg + delm
-    mm     = xmdf + temp
-    argpm  = argpdf - temp
+    mm0     = xmdf + temp
+    argpm0  = argpdf - temp
     t3     = t2 * satellite_batch._t
     t4     = t3 * satellite_batch._t
-    tempa  = tempa - satellite_batch._d2 * t2 - satellite_batch._d3 * t3 - \
+    tempa0  = tempa1 - satellite_batch._d2 * t2 - satellite_batch._d3 * t3 - \
                         satellite_batch._d4 * t4
-    tempe  = tempe + satellite_batch._bstar * satellite_batch._cc5 * (mm.sin() -
+    tempe0  = tempe1 + satellite_batch._bstar * satellite_batch._cc5 * (mm0.sin() -
                         satellite_batch._sinmao)
-    templ  = templ + satellite_batch._t3cof * t3 + t4 * (satellite_batch._t4cof +
+    templ0  = templ1 + satellite_batch._t3cof * t3 + t4 * (satellite_batch._t4cof +
                         satellite_batch._t * satellite_batch._t5cof)
-    # END: ASSUME satellite._isimp IS ALWAYS 0
+
+    mm    = torch.where(satellite_batch._isimp==0,mm0,mm1)
+    argpm = torch.where(satellite_batch._isimp==0,argpm0,argpm1)
+    tempa = torch.where(satellite_batch._isimp==0,tempa0,tempa1)
+    tempe = torch.where(satellite_batch._isimp==0,tempe0,tempe1)
+    templ = torch.where(satellite_batch._isimp==0,templ0,templ1)
 
     nm    = satellite_batch._no_unkozai.clone()
     em    = satellite_batch._ecco.clone()

tests/test_batched_sgp4.py

@@ -44,6 +44,49 @@ def test_sgp4_batched(self):
         #we batch propagate all TLEs at all required times:
         out_batched=dsgp4.propagate_batch(tles_batch,tsinces_batch)
         self.assertTrue(np.allclose(out_non_batched.numpy(),out_batched.numpy()))
+
+    def test_isimp_batched(self):
+        lines = file1.splitlines()
+        index = list(range(1,4,3))
+        tles_batch      = []
+        tsinces_batch    = []
+        out_non_batched = []
+        for i in index:
+            data = []
+            data.append(lines[i])
+            data.append(lines[i+1])
+            data.append(lines[i+2])
+            tle = dsgp4.tle.TLE(data)
+            try: 
+                dsgp4.initialize_tle(tle)
+                if tle._error==0:
+                    tsince = torch.rand(100)*10
+                    tles_batch += [tle]*len(tsince)
+                    tsinces_batch+=[tsince]
+                    out_non_batched+=[dsgp4.propagate(tle,tsince)]
+            except Exception as e:
+                self.assertTrue(str(e).split()==error_string.split())
+        tsinces_batch = torch.cat(tsinces_batch)
+        out_non_batched = torch.cat(out_non_batched)
+        out_batched = dsgp4.propagate_batch(tles_batch,tsinces_batch)
+        self.assertTrue(torch.any(torch.tensor([tle._isimp==1 for tle in tles_batch])))
+        self.assertTrue(np.allclose(out_non_batched.numpy(),out_batched.numpy()))
+
+
+file1 = """
+0 STARLINK-1110
+1 44946U 20001AJ  24101.30511192  .09874184  11821-4  44382-2 0  9991
+2 44946  52.9924  47.4575 0002378 194.8867 193.8851 16.21437839237242
+0 XJS E
+1 45251U 20014C   24100.53376843  .08568485  13520-5  71805-3 0  9994
+2 45251  35.0216  36.8737 0007634 230.7084 292.0228 16.38395432233146
+0 COSMOS 2251 DEB
+1 34427U 93036RU  22068.94647328  .00008100  00000-0  11455-2 0  9999
+2 34427  74.0145 306.8269 0033346  13.0723 347.1308 14.76870515693886
+0 COSMOS 2251 DEB
+1 34428U 93036RV  22068.90158861  .00002627  00000-0  63561-3 0  9999
+2 34428  74.0386 139.1157 0038434 196.7068 279.9147 14.53118052686950
+"""
 
 file="""
 0 COSMOS 2251 DEB