3) Dynamics Along the Tangent

Using Newton’s second law projected along the track tangent: $$ m\frac{dv}{dt}=F_{\text{gravity,tan}}+F_{\text{drag,lin}}+F_{\text{drag,quad}}. $$ With gravity $\mathbf g=[0,-g]^T$, the tangential gravity contribution is $$ F_{\text{gravity,tan}}=m\,\mathbf g\cdot \hat{\mathbf T}(u) =-m g\,T_z(u). $$ A common resistive model is: $$ F_{\text{drag,lin}}=-k_1 v,\qquad F_{\text{drag,quad}}=-k_2 v|v|. $$ So $$ m\frac{dv}{dt} = -m g\,T_z(u)-k_1 v-k_2 v|v|, $$ and $$ \boxed{ \frac{dv}{dt} = -g\,T_z(u)-\frac{k_1}{m}v-\frac{k_2}{m}v|v| }. $$

3) Losses and optimization

The MLP is trained using the two losses with mean squared error (MSE): $$ \mathcal L_{drdu} = \lambda_{drdu}\,\frac{1}{N}\sum_i \left(\hat d_i - d_i\right)^2, $$ $$ \mathcal L_{tz} = \lambda_{tz}\,\frac{1}{N}\sum_i \left(\hat t_i - t_i\right)^2. $$ After comparing plots for $\|dr/du\|$ and $T_z(u)$, we choose $\lambda_{drdu} = 1$ and $\lambda_{tz} = 100$. We optimize with torch.optim.Adam.

3) Initial-condition loss

We could enforce intitial conditions through a loss function $$ \mathcal L_{\text{IC}} = w_{u0}\,\big(u(0)-u_0\big)^2 + w_{v0}\,\big(v(0)-v_0\big)^2. $$ As an elegant alternative we hard constrain these initial conditions into the network using a HardConstrainedPINN wrapper class. It trains parameters $u_{raw}$ and $v_{raw}$ of the original PINN MLP such that $$ \begin{align} u &= u_0 + (1 - exp(-t))\,u_{raw}\\ v &= v_0 + (1 - exp(-t))\,v_{raw}. \end{align} $$

                # 1. HARD-CONSTRAIN THE NETWORK OUTSIDE THE OPTIMIZATION LOOP
                # This forces the network to mathematically satisfy ICs, eliminating ic_loss entirely.
                class HardConstrainedPINN(torch.nn.Module):
                    def __init__(self, raw_pinn, U0, V0):
                        super().__init__()
                        self.pinn = raw_pinn
                        # Register constants to ensure they stay on the correct device
                        self.register_buffer("u0", torch.tensor([[U0]], dtype=torch.float32))
                        self.register_buffer("v0", torch.tensor([[V0]], dtype=torch.float32))

                    def forward(self, t):
                        uv_raw = self.pinn(t)
                        u_raw, v_raw = uv_raw[:, 0:1], uv_raw[:, 1:2]
                        
                        # At t=0, (1 - exp(-t)) is 0, guaranteeing u = u0 and v = v0
                        # As t grows, the network freely learns the residual dynamics
                        u_constrained = self.u0 + (1.0 - torch.exp(-t)) * u_raw
                        v_constrained = self.v0 + (1.0 - torch.exp(-t)) * v_raw
                        
                        return torch.cat([u_constrained, v_constrained], dim=1)
  

5) Training strategy

Unlike traditional ODE solvers, PINNs learn the solution across all time collocation points simultaneously. For better convergence into the correct solution space, we split the time domain $t$ into smaller intervals, training the PINN incrementally with these expanding $t$ intervals. Each interval has a linearly rising distribution generated by a custom function rand_linear_rising and is initially trained with torch.optim.Adam and refined with torch.optim.LBFGS.

                  def rand_linear_rising(shape, a, b):
                    # 1. Generate standard uniform random numbers between 0 and 1
                    u = torch.rand(shape)
                    
                    # 2. Apply the inverse CDF for a linearly rising distribution
                    return a + (b - a) * torch.sqrt(u)
  

Final implementation:

                
                  import os
                  import torch
                  from omegaconf import DictConfig
                  import torch.optim as optim
                  from torch.optim import Adam, lr_scheduler
                  
                  # ---------------- constants ----------------
                  T_MAX = 7
                  MASS = 500.0
                  DRAG_LINEAR = 3.0
                  DRAG_QUADRATIC = 4.0
                  U0 = 0.0
                  V0 = 0.25

                  def load_track_model(ckpt_dir="checkpoints"):
                    track_model = TrackNN()
                    params_path = os.path.join(ckpt_dir, "track_model_output.0.pth")
                    track_model.load_state_dict(torch.load(params_path, map_location="cpu"))
                    return track_model


                  def train():
                      raw_pinn = PINN()
                      pinn = HardConstrainedPINN(raw_pinn, U0, V0)

                      # load frozen track heads
                      track_model = load_track_model()
                      max_epochs = 2000
                      t_int_steps = 12
                      t_max_list = [np.floor(i) for i in np.arange(1,t_int_steps + 1)*T_MAX/t_int_steps]


                      #t_res = torch.linspace(0, T_MAX, 1000).view(-1, 1).requires_grad_(True)
                      track = Track()

                      def compute_loss(current_t_max):
                          # Collocation points where we enforce ODE
                          t_int = rand_linear_rising((256,1), a=0, b=current_t_max)
                          t_res, sort_idx = torch.sort(t_int, dim=0)
                          t_res.requires_grad_(True)

                          # Define custom loss function computation
                          uv = pinn(t_res) 
                          u, v = uv[:, 0:1], uv[:, 1:2]

                          du_dt = torch.autograd.grad(u, t_res, torch.ones_like(u), create_graph=True)[0]
                          dv_dt = torch.autograd.grad(v, t_res, torch.ones_like(v), create_graph=True)[0]

                          # Using clamped u to grab track values safely
                          u_clamped = torch.clamp(u, min=0, max=1)
                          drdu_tz = track_model(u_clamped)
                          drdu, tz = drdu_tz[:, 0:1], drdu_tz[:, 1:2]

                          # ODE residuals
                          res1 = du_dt * drdu - v 
                          res2 = dv_dt + 9.81 * tz + (DRAG_LINEAR / MASS) * v + (DRAG_QUADRATIC / MASS) * v * torch.abs(v)

                          base_res_loss = (1 * (res1**2 + res2**2)).mean()

                          loss = base_res_loss
                          return loss, res1, res2


                      for i, current_t_max in enumerate(t_max_list):
                          print(f"--- t interval {i+1}/{t_int_steps} ---")
                          # --- Phase 1: Adam Pre-training ---
                          print("--- Starting Phase 1: Adam Pre-training ---")
                          opt = Adam(pinn.parameters(), lr=0.001)

                          for epoch in range(max_epochs):
                              opt.zero_grad()
                              loss, res1, res2 = compute_loss(current_t_max)
                              loss.backward()
                              opt.step()
                              
                              if epoch % 500 == 0:
                                  print(f"Adam Step {epoch:05d} | Loss: {loss.item():.2e} | res_u: {(res1**2).mean().item():.2e} | res_v: {(res2**2).mean().item():.2e}")

                          # --- Phase 2: LBFGS Fine-Tuning ---
                          print("--- Starting Phase 2: LBFGS Fine-tuning ---")
                          
                          # Instantiate LBFGS Optimizer with a stable strong_wolfe line search
                          lbfgs_opt = optim.LBFGS(
                              pinn.parameters(),
                              lr=1.0,
                              max_iter=1000,                # Max iterations within the single LBFGS block
                              history_size=100,              # Curvature history size
                              line_search_fn="strong_wolfe" # Modulates learning epoch sizes precisely
                          )
                          # LBFGS requires a self-contained closure function
                          def closure():
                              lbfgs_opt.zero_grad()
                              loss, _, _ = compute_loss(current_t_max)
                              loss.backward()
                              return loss

                          # Execute LBFGS convergence sequence
                          lbfgs_opt.step(closure)
                          
                          # Print the final loss metrics
                          final_loss = closure()