Issue batching with tensors

[pearsonkyle]

Hello,

I'm running into some issues when trying to use this model with batches of TLEs. Below is a simplified implementation of the code giving me issues.

import dsgp4
import dataclasses
import torch
import torch.nn as nn
import torch.nn.functional as F

torch.set_default_dtype(torch.float32)
DEFAULT_DEVICE = "cuda:7"

#let's first define the Earth radius according to WSG-84:
r_earth=dsgp4.util.get_gravity_constants('wgs-84')[2].numpy() # km
v_earth=10 # km/s

@dataclasses.dataclass
class OrbitModel:
    satellite_catalog_number: int
    classification: str
    international_designator: str
    epoch_year: int
    epoch_days: float
    ephemeris_type: int
    element_number: int
    revolution_number_at_epoch: int
    mean_motion: float
    mean_motion_first_derivative: float
    mean_motion_second_derivative: float
    eccentricity: float
    inclination: float
    argument_of_perigee: float
    raan: float
    mean_anomaly: float
    b_star: float

    def _create_tle(self, initialize=True, **kwargs):
        tle_data = {
            'satellite_catalog_number': self.satellite_catalog_number,
            'classification': self.classification,
            'international_designator': self.international_designator,
            'epoch_year': self.epoch_year,
            'epoch_days': self.epoch_days,
            'ephemeris_type': self.ephemeris_type,
            'element_number': self.element_number,
            'revolution_number_at_epoch': self.revolution_number_at_epoch,
            'mean_motion': self.mean_motion,
            'mean_motion_first_derivative': self.mean_motion_first_derivative,
            'mean_motion_second_derivative': self.mean_motion_second_derivative,
            'eccentricity': self.eccentricity,
            'inclination': self.inclination,
            'argument_of_perigee': self.argument_of_perigee,
            'raan': self.raan,
            'mean_anomaly': self.mean_anomaly,
            'b_star': self.b_star
        }
        tle_data.update(kwargs)
        tle = dsgp4.tle.TLE(tle_data)
        if initialize:
            dsgp4.initialize_tle(tle)
        return tle

    def __post_init__(self):
        self.tle = self._create_tle()

    def update(self, **kwargs):
        for key, value in kwargs.items():
            setattr(self, key, value)
        self.tle = self._create_tle()

    def propagate(self, tsinces):
        state_teme = dsgp4.propagate(self.tle, tsinces)
        return state_teme[:,0], state_teme[:,1] # pos, vel

class OrbitMLP(nn.Module):
    def __init__(self, input_size, sequence_size, model, bounds, dropout_rate=0.1):
        super(OrbitMLP, self).__init__()
        self.fc1 = nn.Linear(input_size, 16)
        self.fc2 = nn.Linear(16, 16)
        self.fc3 = nn.Linear(16, len(bounds))
        self.dropout = nn.Dropout(dropout_rate)
        self.output_activation = nn.Sigmoid()

        self.model = model
        self.bounds = bounds
        self.param_keys = [k for k in bounds]
        self.tsinces = torch.arange(sequence_size, dtype=torch.float32, requires_grad=True)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        output = self.output_activation(self.fc3(x))[0]

        # create set of TLEs based on output
        tle_batch = self._create_tle_batch(output)

        return tle_batch.unsqueeze(0)

    # create new TLEs and propagate orbits
    def _create_tle_batch(self, output):
        tles = []
        sv = torch.zeros((output.shape[0], 6)).to(DEFAULT_DEVICE)
        for i in range(output.shape[0]):
            
            # create new dictionary of parameters
            new_args = {k: output[i][j] for j, k in enumerate(self.param_keys)}
            
            # invert the normalization to put back in normal units
            iargs = self.invert_normalize_params(new_args)
            
            # create a new tle
            tle = self.model._create_tle(**iargs)
            tles.append(tle)

        _, tle_batch = dsgp4.initialize_tle(tles, with_grad=True)
        return dsgp4.propagate_batch(tle_batch, self.tsinces)

    def invert_normalize_params(self, nparams):
        params = {}
        for param, (min_val, max_val) in self.bounds.items():
            params[param] = nparams[param] * (max_val - min_val) + min_val
        return params


dsgp4_model = OrbitModel(
    satellite_catalog_number=43437,
    classification='U',
    international_designator='18100A',
    epoch_year=2020,
    epoch_days=143.90384230,
    ephemeris_type=0,
    element_number=9996,
    revolution_number_at_epoch=0,
    mean_motion=0.0011,
    mean_motion_first_derivative=0.,
    mean_motion_second_derivative=0.,
    eccentricity=0.02,
    inclination=1.7074,
    argument_of_perigee=2.15,
    raan=4.3618,
    mean_anomaly=4.5224,
    b_star=0.000001
)

# Extract the time and position data
t = torch.arange(0,720, dtype=torch.float32)
pos,vel = dsgp4_model.propagate(t)


# Initialize the model
bounds={
    'mean_motion': (0.001, 0.00115),
    'eccentricity': (0.0, 0.0001),
    'inclination': (0, 6.28),
    'argument_of_perigee': (-6.28, 6.28), #mostly for eccentric orbits
    'raan': (3.14, 6.28),
    'mean_anomaly': (0, 3.14)
}

omlp = OrbitMLP(6, len(pos), dsgp4_model, bounds)

# create one sample of size: n_batch, n_time, n_feature
X = torch.hstack([pos,vel]).unsqueeze(0)

omlp(X).shape

The line that throws an error is: _, tle_batch = dsgp4.initialize_tle(tles, with_grad=True) with the message, Warning: 1 TLEs were not initialized because they are in deep space. Deep space propagation is currently not supported. even though I'm able to evaluate it prior to the MLP. Does it have something to do with the input tensors having gradients? Even though the TLE is staying constant, I'm trying to have the network predict the value at each time step in a very simple example.

python 3.10.12
torch 2.2.0

[Sceki]

Hi @pearsonkyle ,

Thanks a lot for the report (for next times, please try to put a minimal code to reproduce the error, it took me a while :D ).

Anyway, it seems that you are updating the TLE using as elements of the dictionary torch tensors with leaf (in this line --> new_args = {k: output[i][j] for j, k in enumerate(self.param_keys)}).

This triggers an error internally because at some point the dictionary elements are deepcopied with copy.deepcopy (since they are expected to be int or float) and torch tensors with leaf do not allow that.

You might want to use floats to update the dictionary of tle data. Try for instance something like --> new_args = {k: output[i][j].item() for j, k in enumerate(self.param_keys)} instead. And, please, let me know how it goes!

PS: In any case, I opened #15 because the warning being thrown is currently not helpful for debugging (the actual error is instead related to the deepcopy as I said (something like: RuntimeError: Only Tensors created explicitly by the user (graph leaves) support the deepcopy protocol at the moment.). I also added the possibility to use tensors with leaf in the dictionary when copying the TLE, so even your initial setup should not trigger anymore, although I think it's best to use floats there if possible.

[Sceki]

Closing this as the issue should have been resolved, and now the new version is available both in pip and conda, but please let me know if there are any issues and we can still reopen

[pearsonkyle]

Many thanks for the quick reply @Sceki , I will test it again with the latest version.

I had a chance to test the suggestion and while it gets the code to run it doesn't propagate the gradient from the MLP output into the SGP4 model. Unfortunately, the .item strips the value of it's data type which should be a tensor with a gradient property. The loop can run but it doesn't actually regress or learn anything without the gradient

[Sceki]

HI @pearsonkyle .. this is not a bug but a matter of implementation :)

I still have not uploaded the tutorial on how to do ML-dSGP4 training, but it will be published soon and it will clarify some of your doubts.

In any case, I think perhaps to get some inspiration on how to do this, you can have a look at the tutorial on Gradient Based Optimization (https://esa.github.io/dSGP4/notebooks/gradient_based_optimization.html), and the code (e.g. https://github.com/esa/dSGP4/blob/v1.0.2/dsgp4/newton_method.py#L113).

I hope these help for now to get started and to have a way to keep track of the gradients while updating the TLE and propagating.