From 8a0efc0090a402d3a7458677b7af961aaeaf7291 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Wed, 20 May 2026 20:51:58 +0200
Subject: [PATCH 01/31] tunnel factor added

---
 ocp.py                      |  6 ++--
 simulations/sim_launcher.py | 36 ++++++++++++++++++++
 track.py                    | 68 ++++++++++++++++++++++++++++++++++---
 train.py                    | 47 +++++++++++++++----------
 utils.py                    | 35 ++++++++++++++++++-
 5 files changed, 167 insertions(+), 25 deletions(-)
 create mode 100644 simulations/sim_launcher.py

diff --git a/ocp.py b/ocp.py
index 3ca7985..76ce2ce 100644
--- a/ocp.py
+++ b/ocp.py
@@ -6,7 +6,7 @@
 
 from track import computeDiscretizationPoints
 
-from utils import Options, var, postProcessDataFrame, splitLosses
+from utils import Options, var, postProcessDataFrame, splitLosses, computeTunnelFactor
 
 
 class OptionsCasadiSolver(Options):
@@ -194,9 +194,11 @@ def __init__(self, train, track, optsDict={}):
                 # gradient and curvature of current index
                 grad = self.points.iloc[i]['Gradient [permil]']/1e3
                 curv = self.points.iloc[i]['Curvature [1/m]']
+                crossSection = self.points.iloc[i]['CrossSection [m^2]']
+                tunnelFactor = computeTunnelFactor(crossSection, train)
 
                 # acceleration constraints
-                g += [trainModel.accelerationFun(ca.vertcat(time[i], velSq[i]), ca.vcat(u), grad, curv)]
+                g += [trainModel.accelerationFun(ca.vertcat(time[i], velSq[i]), ca.vcat(u), grad, curv, tunnelFactor)]
                 lbg += [accMin]
                 ubg += [accMax]
 
diff --git a/simulations/sim_launcher.py b/simulations/sim_launcher.py
new file mode 100644
index 0000000..c58d9dd
--- /dev/null
+++ b/simulations/sim_launcher.py
@@ -0,0 +1,36 @@
+from ocp import casadiSolver
+
+if __name__ == '__main__':
+
+    from train import Train
+    from track import Track
+
+    # Timetable
+    startPosition = 0       # [m]
+    endPosition = 20000     # [m]
+    duration = 60*20        # [s]
+
+    train = Train(config={'id':'SBB_Flirt_2'}, pathJSON='../trains')
+
+    track = Track(config={'id':'CH_ZH_LU'}, pathJSON='../tracks')
+    # track = Track(config={'id':'CH_StGallen_Wil'}, pathJSON='../tracks')
+    track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
+
+    opts = {'numIntervals':200, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':1}, 'energyOptimal':True}
+
+    solver = casadiSolver(train, track, opts)
+
+    df, stats = solver.solve(duration)
+
+    # print some info
+    if df is not None:
+
+        print("")
+        print("Objective value = {:.2f} {}".format(stats['Cost'], 'kWh' if solver.opts.energyOptimal else 's'))
+        print("")
+        print("Maximum acceleration: {:5.2f}, with bound {}".format(df.max()['Acceleration [m/s^2]'], train.accMax if train.accMax is not None else 'None'))
+        print("Maximum deceleration: {:5.2f}, with bound {}".format(df.min()['Acceleration [m/s^2]'], train.accMin if train.accMin is not None else 'None'))
+
+    else:
+
+        print("Solver failed!")
\ No newline at end of file
diff --git a/track.py b/track.py
index c2ad35c..e0b0d4f 100644
--- a/track.py
+++ b/track.py
@@ -1,12 +1,12 @@
 import json
-import math
 import os
 import sys
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 
-from utils import checkTTOBenchVersion, convertUnit
+from utils import checkTTOBenchVersion, convertUnit, pickEquallySpacedPoints
+
 
 def importTuples(tuples, xLabel, yLabels):
     """
@@ -95,7 +95,7 @@ def computeDiscretizationPoints(track, numIntervals):
 
     df1 = track.mergeDataFrames()
 
-    pos = np.linspace(0, track.length, numIntervals + 1 - (len(df1) - 1))
+    pos = pickEquallySpacedPoints(0, track.length, numIntervals, df1.index.to_numpy(dtype=float))
 
     df2 = pd.DataFrame({'position [m]':pos}).set_index('position [m]')
     df3 = df2.join(df1, how='outer').ffill()
@@ -149,6 +149,11 @@ def __init__(self, config, pathJSON='tracks'):
                                    data['curvatures']['units']['radius at end'] if 'curvatures' in data else "m",
                                    config['clothoidSamplingInterval'] if 'clothoidSamplingInterval' in config else None)
 
+        self.importTunnelTuples(data['tunnels']['values'] if 'tunnels' in data else [(0.0, 0.0, "infinity")],
+                                data['tunnels']['units']['length'] if 'tunnels' in data else 'm',
+                                data['tunnels']['units']['cross_section'] if 'tunnels' in data else 'm^2')
+
+
         numStops = len(data['stops']['values'])
         indxDeparture = config['from'] if 'from' in config else 0
         indxDestination = config['to'] if 'to' in config else numStops-1
@@ -193,6 +198,11 @@ def curvaturesOk(self):
         return True if self.curvatures.shape[0] > 0 and checkDataFrame(self.curvatures, self.length) else False
 
 
+    def crossSectionsOk(self):
+
+        return True if self.crossSections.shape[0] > 0 and checkDataFrame(self.crossSections, self.length) else False
+
+
     def checkFields(self):
 
         if not self.lengthOk():
@@ -215,6 +225,10 @@ def checkFields(self):
 
             raise ValueError("Issue with track curvatures!")
 
+        if not self.crossSectionsOk():
+
+            raise ValueError("Issue with track cross sections!")
+
 
     def importGradientTuples(self, tuples, unit='permil'):
 
@@ -348,6 +362,49 @@ def sampleClothoid(self, tuples, ds=None):
         return result
 
 
+    def importTunnelTuples(self, tuples, unitLength='m', unitCrossSection='m^2'):
+
+        if not self.lengthOk():
+
+            raise ValueError("Cannot import tunnels without a valid track length!")
+
+        if unitLength not in {'m', 'km'} or unitCrossSection not in {'m^2'}:
+
+            raise ValueError("Specified tunnel units not supported!")
+
+        tuples = [(p, convertUnit(l, unitLength), convertUnit(c, unitCrossSection)) for p,l,c in tuples]
+        self.tunnels = importTuples(tuples, 'Position [m]', ['Length [m]', 'CrossSection [m^2]'])
+
+
+        # get end of tunnel positions and assign them a cross section of inf
+        positions = self.tunnels.index.astype(float)
+        tunnelLengths = self.tunnels["Length [m]"].astype(float)
+
+        endOfTunnelPositions = positions + tunnelLengths
+
+        # a tunnel may change its cross section, therefore some end of tunnel positions need to be removed
+        endOfTunnelPositions = [
+            e for e in endOfTunnelPositions
+            if not any(abs((p - e)) < 0.1 for p in positions)
+        ]
+
+        openTrack_df = pd.DataFrame({"Length [m]": 0.0, "CrossSection [m^2]": float("inf")}, index=endOfTunnelPositions)
+        openTrack_df.index.name = self.tunnels.index.name
+        self.tunnels = pd.concat([self.tunnels, openTrack_df]).sort_index()
+
+        if positions[0] != 0.0:
+            first_row = {"Position [m]": 0.0, "Length [m]": 0.0, "CrossSection [m^2]": float("inf")}
+            self.tunnels.loc[0] = first_row
+            self.tunnels = self.tunnels.sort_index()
+
+        self.tunnels.drop(columns=["Length [m]"], inplace=True)
+
+        self.crossSections = self.tunnels
+        del self.tunnels
+
+        checkDataFrame(self.crossSections, self.length)
+
+
     def reverse(self):
         # switch to opposite trip
 
@@ -368,6 +425,7 @@ def flipData(df):
         self.gradients = -flipData(self.gradients)
         self.speedLimits = flipData(self.speedLimits)
         self.curvatures = -flipData(self.curvatures)
+        self.crossSections = flipData(self.crossSections)
 
         self.title = self.title + ' (reversed)'
 
@@ -380,7 +438,8 @@ def mergeDataFrames(self):
         """
 
         joinedGradientsAndSpeedLimits = self.gradients.join(self.speedLimits, how='outer').fillna(method='ffill')
-        return self.curvatures.join(joinedGradientsAndSpeedLimits, how='outer').fillna(method='ffill')
+        joined = self.curvatures.join(joinedGradientsAndSpeedLimits, how='outer').fillna(method='ffill')
+        return self.crossSections.join(joined, how='outer').fillna(method='ffill')
 
     def print(self):
         """
@@ -448,6 +507,7 @@ def crop(dfIn):
         self.speedLimits = crop(self.speedLimits)
         self.gradients = crop(self.gradients)
         self.curvatures = crop(self.curvatures)
+        self.crossSections = crop(self.crossSections)
 
 
 if __name__ == '__main__':
diff --git a/train.py b/train.py
index dd8567b..43f1a7b 100644
--- a/train.py
+++ b/train.py
@@ -95,6 +95,10 @@ def __init__(self, config, pathJSON='trains') -> None:
 
         self.r2 = convertUnit(data['rolling resistance r2']['value'], data['rolling resistance r2']['unit'])  # quadratic term [N/(m/s)^2]
 
+        self.t_24 = convertUnit(data['tunnel resistance 24 m^2']['value'], data['tunnel resistance 24 m^2']['unit'])  # [kg/m]
+
+        self.t_40 = convertUnit(data['tunnel resistance 40 m^2']['value'], data['tunnel resistance 40 m^2']['unit'])  # [kg/m]
+
         # TODO: unify with case of dynamic efficiency
         if 'efficiency traction' in data or 'efficiency reg brake' in data:
 
@@ -225,32 +229,39 @@ def __init__(self, sr0, sr1, sr2, rho=1, g=9.81, withPnBrake=True) -> None:
 
         # states
 
-        time = ca.MX.sym('time')
-        velocitySquared = ca.MX.sym('velocitySquared')
+        time = ca.MX.sym('time')                            # [s]
+        velocitySquared = ca.MX.sym('velocitySquared')      # [m^2/s^2]
 
         x = ca.vertcat(time, velocitySquared)
 
         # controls
 
-        traction = ca.MX.sym('traction')
-        pnBrake = ca.MX.sym('pnBrake')
+        traction = ca.MX.sym('traction')                    # [N/kg]
+        pnBrake = ca.MX.sym('pnBrake')                      # [N/kg]
 
         u = ca.vertcat(traction, pnBrake if withPnBrake else [])
 
         # parameters
 
-        gradient = ca.MX.sym('gradient')
-        curvature = ca.MX.sym('curvature')
+        gradient = ca.MX.sym('gradient')                    # [-]     -> values between [0,0.2]
+        curvature = ca.MX.sym('curvature')                  # [1/m]   -> values between [0,0.004]
+        tunnelFactor = ca.MX.sym('tunnelFactor')            # [1/m]
         ds = ca.MX.sym('ds')
 
-        p = ca.vertcat(gradient, curvature, ds)
+        p = ca.vertcat(gradient, curvature, tunnelFactor, ds)
 
         # ODE
 
-        rollingResistance = sr0 + sr1*ca.sqrt(velocitySquared) + sr2*velocitySquared
-        curvatureResistance = ca.if_else(ca.fabs(curvature)<=1/300, g*0.5*ca.fabs(curvature)/(1-30*ca.fabs(curvature)),
-                                         g*0.65*ca.fabs(curvature)/(1-55*ca.fabs(curvature)))
-        acceleration = traction + (pnBrake if withPnBrake else 0) - rollingResistance - g*gradient*(1/rho) - curvatureResistance*(1/rho)
+        rollingResistance = sr0 + sr1*ca.sqrt(velocitySquared) + sr2*velocitySquared                # [N/kg]
+        gradientResistance = g*gradient*(1/rho)                                                     # [N/kg]
+        curvatureResistance = ca.if_else(ca.fabs(curvature)<=1/300,
+                                            g*0.5*ca.fabs(curvature)/(1-30*ca.fabs(curvature)),
+                                            g*0.65*ca.fabs(curvature)/(1-55*ca.fabs(curvature))
+                                         )                                                          # [N/kg]
+        tunnelResistance = tunnelFactor * velocitySquared                                           # [N/kg]
+
+        acceleration = traction + (pnBrake if withPnBrake else 0) - rollingResistance - gradientResistance - curvatureResistance - tunnelResistance # [m/s^2]
+
         timeODE = 1/ca.sqrt(velocitySquared)
         velocityODE = 2*acceleration
 
@@ -263,7 +274,7 @@ def __init__(self, sr0, sr1, sr2, rho=1, g=9.81, withPnBrake=True) -> None:
 
         self.ode = fExplicit
         self.acceleration = acceleration
-        self.accelerationFun = ca.Function('a', [x, u, gradient, curvature], [acceleration])
+        self.accelerationFun = ca.Function('a', [x, u, gradient, curvature, tunnelFactor], [acceleration])
         self.rollingResistance = rollingResistance
         self.parameters = p
         self.controls = u
@@ -343,7 +354,7 @@ def __init__(self, model, solver, optsDict={}) -> None:
             self.eval = ca.Function('xNxt', [model.states, ca.vertcat(model.controls, model.parameters), ca.MX.sym('ds')], [eval])
 
 
-    def solve(self, time, velocitySquared, ds, traction=0, pnBrake=0, gradient=0, curvature=0):
+    def solve(self, time, velocitySquared, ds, traction=0, pnBrake=0, gradient=0, curvature=0, tunnelFactor=0):
 
         withPnBrake = self.model.withPnBrake
 
@@ -353,7 +364,7 @@ def solve(self, time, velocitySquared, ds, traction=0, pnBrake=0, gradient=0, cu
 
         x0 = ca.vertcat(time, velocitySquared)
         u0 = ca.vertcat(traction, pnBrake if withPnBrake else [])
-        p0 = ca.vertcat(gradient, curvature, ds)
+        p0 = ca.vertcat(gradient, curvature, tunnelFactor, ds)
         x1 = self.eval(x0, ca.vertcat(u0, p0), 1)
 
         out = {}
@@ -401,11 +412,11 @@ def initLosses(self, lossesTrFun, lossesRgbFun, totalMass, solver='CVODES'):
             raise ValueError("Unknown solver!")
 
 
-    def calcLosses(self, velocity, dt, traction=0, pnBrake=0, gradient=0, curvature=0):
+    def calcLosses(self, velocity, dt, traction=0, pnBrake=0, gradient=0, curvature=0, tunnelFactor=0):
 
         mdl = self.model
 
-        out = self.lossesIntegrator(ca.vertcat(velocity, 0, 0), ca.vertcat(traction, pnBrake if mdl.withPnBrake else [], gradient, curvature), dt)
+        out = self.lossesIntegrator(ca.vertcat(velocity, 0, 0), ca.vertcat(traction, pnBrake if mdl.withPnBrake else [], gradient, curvature, tunnelFactor), dt)
 
         lossesTr, lossesRgb = out[1], out[2]
 
@@ -441,12 +452,12 @@ def initRollingResistance(self, solver='CVODES'):
             raise ValueError("Unknown solver!")
 
 
-    def calcRollingResistance(self, velocity, ds, traction=0, pnBrake=0, gradient=0, curvature=0):
+    def calcRollingResistance(self, velocity, ds, traction=0, pnBrake=0, gradient=0, curvature=0, tunnelFactor=0):
 
         mdl = self.model
 
         out = self.rollingResistanceIntegrator(ca.vertcat(velocity**2, 0), ca.vertcat(traction, pnBrake if mdl.withPnBrake else [], gradient,
-                                                                                      curvature, ds), 1)
+                                                                                      curvature, tunnelFactor, ds), 1)
 
         losses = out[1]
 
diff --git a/utils.py b/utils.py
index fb30e42..bf38dee 100644
--- a/utils.py
+++ b/utils.py
@@ -369,7 +369,7 @@ def convertUnit(value, unit):
     Convert from any known unit to internally used unit.
     """
 
-    if unit in {'m', 'm/s', 'permil', 'kg', 'W', 'N', 'm/s^2', '-', 'N/(m/s)', 'N/(m/s)^2', 'kg/m'}:
+    if unit in {'m', 'm/s', 'm^2', 'permil', 'kg', 'W', 'N', 'm/s^2', '-', 'N/(m/s)', 'N/(m/s)^2', 'kg/m'}:
 
         valueOut = value
 
@@ -479,6 +479,39 @@ def latexify():
     return latexFound
 
 
+def computeTunnelFactor(cross_section, train):
+
+    if cross_section == 0:
+        return 0 # no tunnel
+
+    total_mass = train.mass * train.rho
+    t_24 = train.t_24
+    t_40 = train.t_40
+
+    return t_24/total_mass if cross_section < 30 else t_40/total_mass
+
+
+def pickEquallySpacedPoints(startPoint, endPoint, numIntervals, requiredPoints):
+
+    np.random.seed(42)
+
+    if len(requiredPoints) > numIntervals + 1:
+        raise ValueError(f"Too many required points ({len(requiredPoints)}) for N.")
+
+    num_of_remaining_points = numIntervals + 1 - len(requiredPoints)
+    m = 1 # number of points to oversample
+
+    while True:
+        cand = np.linspace(startPoint, endPoint, num_of_remaining_points + m + 2)[1:-1]  # oversample to avoid overlaps
+        cand = np.round(cand, 0)
+        out = np.unique(np.r_[requiredPoints, cand])
+        cand_without_required = out[~np.isin(out, requiredPoints)]
+        if len(cand_without_required) >= num_of_remaining_points:
+            picked_points = np.random.choice(cand_without_required, size=num_of_remaining_points, replace=False)
+            return picked_points
+        m *= 2
+
+
 if __name__ == '__main__':
 
     pass

From 379d22128b6a959230b5eb0a6011a580d73669ca Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Thu, 21 May 2026 10:56:00 +0200
Subject: [PATCH 02/31] speed limits train length adjusted

---
 simulations/sim_launcher.py |  1 +
 track.py                    | 44 ++++++++++++++++++++++++++++++++++++-
 train.py                    |  2 ++
 utils.py                    | 30 +++++++++++++++++++++++++
 4 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/simulations/sim_launcher.py b/simulations/sim_launcher.py
index c58d9dd..36ec6d5 100644
--- a/simulations/sim_launcher.py
+++ b/simulations/sim_launcher.py
@@ -15,6 +15,7 @@
     track = Track(config={'id':'CH_ZH_LU'}, pathJSON='../tracks')
     # track = Track(config={'id':'CH_StGallen_Wil'}, pathJSON='../tracks')
     track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
+    track.updateTrainLengthDependentValues(train)
 
     opts = {'numIntervals':200, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':1}, 'energyOptimal':True}
 
diff --git a/track.py b/track.py
index e0b0d4f..ba535fe 100644
--- a/track.py
+++ b/track.py
@@ -5,7 +5,7 @@
 import numpy as np
 import pandas as pd
 
-from utils import checkTTOBenchVersion, convertUnit, pickEquallySpacedPoints
+from utils import checkTTOBenchVersion, convertUnit, pickEquallySpacedPoints, plotSpeedLimits
 
 
 def importTuples(tuples, xLabel, yLabels):
@@ -510,6 +510,48 @@ def crop(dfIn):
         self.crossSections = crop(self.crossSections)
 
 
+    def updateTrainLengthDependentValues(self, train):
+
+        self.updateSpeedLimitsToTrainLength(train.length)
+
+
+    def updateSpeedLimitsToTrainLength(self, trainLength):
+
+        v = self.speedLimits["Speed limit [m/s]"].to_numpy(dtype=float)
+        pos = self.speedLimits.index.to_numpy(dtype=float)
+
+        if len(pos) > 1:
+
+            pos_adj = []
+            v_adj = []
+
+            for i in range(len(pos)):
+                new_pos = pos[i]
+
+                # Delay speed increases by train length
+                if i > 0 and v[i] > v[i - 1]:
+                    new_pos += trainLength
+
+                # Skip points outside the track
+                if new_pos >= self.length:
+                    continue
+
+                # Remove previous points that are now after this point
+                while pos_adj and pos_adj[-1] > new_pos:
+                    pos_adj.pop()
+                    v_adj.pop()
+
+                pos_adj.append(new_pos)
+                v_adj.append(v[i])
+
+            plotSpeedLimits(self, np.asarray(pos_adj, dtype=float), np.asarray(v_adj, dtype=float))
+
+            self.speedLimits = pd.DataFrame(
+                {"Speed limit [m/s]": v_adj},
+                index=pos_adj,
+            )
+
+
 if __name__ == '__main__':
 
     # Example on how to load and plot a track
diff --git a/train.py b/train.py
index 43f1a7b..209ba48 100644
--- a/train.py
+++ b/train.py
@@ -65,6 +65,8 @@ def __init__(self, config, pathJSON='trains') -> None:
             raise ValueError("Redundant fields in train configuration: {}!".format(', '.join(set(config) - usedFields)))
 
         # read data
+        self.length = convertUnit(data['length']['value'], data['length']['unit'])  # train length [m]
+
         self.mass = convertUnit(data['mass']['value'], data['mass']['unit'])  # train mass [kg]
 
         self.rho = convertUnit(data['rho']['value'], data['rho']['unit'])  # rotating-mass factor [-]
diff --git a/utils.py b/utils.py
index bf38dee..749c5a5 100644
--- a/utils.py
+++ b/utils.py
@@ -7,6 +7,9 @@
 from types import MethodType
 from distutils.spawn import find_executable
 
+from matplotlib import pyplot as plt
+
+
 def var(tag, dim=None):
     "Wrapper to create symbolic variables in casadi."
 
@@ -512,6 +515,33 @@ def pickEquallySpacedPoints(startPoint, endPoint, numIntervals, requiredPoints):
         m *= 2
 
 
+def plotSpeedLimits(track, pos_adj, v_adj):
+
+    v_limits = track.speedLimits["Speed limit [m/s]"].to_numpy(dtype=float)
+    pos_v_limits = track.speedLimits.index.to_numpy(dtype=float)
+
+    v_limits = np.append(v_limits, v_limits[-1])
+    pos_v_limits = np.append(pos_v_limits, track.length)
+
+    v_adj = np.append(v_adj, v_adj[-1])
+    pos_adj = np.append(pos_adj, track.length)
+
+    fig, ax = plt.subplots(figsize=(16, 8))
+
+    ax.step(pos_v_limits / 1000, v_limits  * 3.6, where='post', label="piecewise constant speed limit")
+    ax.step(pos_adj/1000, v_adj * 3.6, where='post', linestyle="--", label="train length adjusted speed limit")
+
+    ax.set_title("Speed Limits")
+    ax.set_xlabel('Position [km]')
+    ax.set_ylabel('Velocity [km/h]')
+    ax.grid(True, which='both', linestyle='--', alpha=0.5)
+    ax.legend(loc="upper right")
+    ax.set_xlim(0, track.length / 1000)
+    ax.figure.tight_layout()
+
+    plt.show()
+
+
 if __name__ == '__main__':
 
     pass

From c79639e5bb68a1172409f53d14423a9af10e8432 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Thu, 21 May 2026 14:40:19 +0200
Subject: [PATCH 03/31] wip

---
 ocp.py                      |  3 ++-
 simulations/sim_launcher.py |  2 +-
 track.py                    | 48 ++++++++++++++++++++++++++++++++++++-
 train.py                    | 28 ++++++++++++----------
 utils.py                    | 24 +++++++++++++++++++
 5 files changed, 90 insertions(+), 15 deletions(-)

diff --git a/ocp.py b/ocp.py
index 76ce2ce..658f264 100644
--- a/ocp.py
+++ b/ocp.py
@@ -193,12 +193,13 @@ def __init__(self, train, track, optsDict={}):
 
                 # gradient and curvature of current index
                 grad = self.points.iloc[i]['Gradient [permil]']/1e3
+                gradLinearTerm = self.points.iloc[i]["Gradient linear term [permil]"]/1e3
                 curv = self.points.iloc[i]['Curvature [1/m]']
                 crossSection = self.points.iloc[i]['CrossSection [m^2]']
                 tunnelFactor = computeTunnelFactor(crossSection, train)
 
                 # acceleration constraints
-                g += [trainModel.accelerationFun(ca.vertcat(time[i], velSq[i]), ca.vcat(u), grad, curv, tunnelFactor)]
+                g += [trainModel.accelerationFun(ca.vertcat(time[i], velSq[i]), ca.vcat(u), grad, gradLinearTerm, curv, tunnelFactor)]
                 lbg += [accMin]
                 ubg += [accMax]
 
diff --git a/simulations/sim_launcher.py b/simulations/sim_launcher.py
index 36ec6d5..832e640 100644
--- a/simulations/sim_launcher.py
+++ b/simulations/sim_launcher.py
@@ -17,7 +17,7 @@
     track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
     track.updateTrainLengthDependentValues(train)
 
-    opts = {'numIntervals':200, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':1}, 'energyOptimal':True}
+    opts = {'numIntervals':400, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':1}, 'energyOptimal':True}
 
     solver = casadiSolver(train, track, opts)
 
diff --git a/track.py b/track.py
index ba535fe..3387420 100644
--- a/track.py
+++ b/track.py
@@ -5,7 +5,7 @@
 import numpy as np
 import pandas as pd
 
-from utils import checkTTOBenchVersion, convertUnit, pickEquallySpacedPoints, plotSpeedLimits
+from utils import checkTTOBenchVersion, convertUnit, pickEquallySpacedPoints, plotSpeedLimits, plotGradients
 
 
 def importTuples(tuples, xLabel, yLabels):
@@ -513,6 +513,7 @@ def crop(dfIn):
     def updateTrainLengthDependentValues(self, train):
 
         self.updateSpeedLimitsToTrainLength(train.length)
+        self.updateGradientsToTrainLength(train.length)
 
 
     def updateSpeedLimitsToTrainLength(self, trainLength):
@@ -552,6 +553,51 @@ def updateSpeedLimitsToTrainLength(self, trainLength):
             )
 
 
+    def updateGradientsToTrainLength(self, trainLength):
+
+        g = self.gradients["Gradient [permil]"].to_numpy(dtype=float)
+        pos = self.gradients.index.to_numpy(dtype=float)
+        slopes = np.r_[0,(g[1:]-g[:-1])/trainLength]
+
+        if len(pos) > 1:
+
+            pos_adj = np.sort(np.r_[pos, pos + trainLength])
+
+            pos_adj = pos_adj[pos_adj < self.length]
+
+            g_adj = [g[0]]
+            g_linear = [0]
+
+            for idx in range(1,len(pos_adj)):
+
+                currentPosition = pos_adj[idx]
+                previousPosition = pos_adj[idx - 1]
+
+                currentGradient = g_adj[idx-1] + (currentPosition-previousPosition)*g_linear[idx-1]
+
+                epsilon = 0.001
+                list_indices = []
+
+                for idx2 in range(len(pos)-1):
+
+                    if currentPosition - trainLength + epsilon < pos[idx2] < currentPosition + epsilon:
+                        list_indices.append(idx2)
+
+                currentLinearTerm = sum(slopes[list_indices])
+
+                g_adj.append(currentGradient)
+                g_linear.append(currentLinearTerm)
+
+            plotGradients(self, np.asarray(pos_adj, dtype=float), np.asarray(g_adj, dtype=float), np.asarray(g_linear, dtype=float))
+
+            self.gradients = pd.DataFrame(
+                {"Gradient [permil]": g_adj, "Gradient linear term [permil]": g_linear},
+                index=pos_adj,
+            )
+
+
+
+
 if __name__ == '__main__':
 
     # Example on how to load and plot a track
diff --git a/train.py b/train.py
index 209ba48..ec7c417 100644
--- a/train.py
+++ b/train.py
@@ -233,8 +233,9 @@ def __init__(self, sr0, sr1, sr2, rho=1, g=9.81, withPnBrake=True) -> None:
 
         time = ca.MX.sym('time')                            # [s]
         velocitySquared = ca.MX.sym('velocitySquared')      # [m^2/s^2]
+        position = ca.MX.sym('position')                    # [m]
 
-        x = ca.vertcat(time, velocitySquared)
+        x = ca.vertcat(time, velocitySquared, position)
 
         # controls
 
@@ -245,17 +246,18 @@ def __init__(self, sr0, sr1, sr2, rho=1, g=9.81, withPnBrake=True) -> None:
 
         # parameters
 
-        gradient = ca.MX.sym('gradient')                    # [-]     -> values between [0,0.2]
-        curvature = ca.MX.sym('curvature')                  # [1/m]   -> values between [0,0.004]
-        tunnelFactor = ca.MX.sym('tunnelFactor')            # [1/m]
+        gradient = ca.MX.sym('gradient')                        # [-]     -> values between [0,0.2]
+        gradientLinearTerm = ca.MX.sym('gradientLinearTerm')    # [-]
+        curvature = ca.MX.sym('curvature')                      # [1/m]   -> values between [0,0.004]
+        tunnelFactor = ca.MX.sym('tunnelFactor')                # [1/m]
         ds = ca.MX.sym('ds')
 
-        p = ca.vertcat(gradient, curvature, tunnelFactor, ds)
+        p = ca.vertcat(gradient, gradientLinearTerm, curvature, tunnelFactor, ds)
 
         # ODE
 
         rollingResistance = sr0 + sr1*ca.sqrt(velocitySquared) + sr2*velocitySquared                # [N/kg]
-        gradientResistance = g*gradient*(1/rho)                                                     # [N/kg]
+        gradientResistance = g*(1/rho)*(gradient+gradientLinearTerm*position)                                                  # [N/kg]
         curvatureResistance = ca.if_else(ca.fabs(curvature)<=1/300,
                                             g*0.5*ca.fabs(curvature)/(1-30*ca.fabs(curvature)),
                                             g*0.65*ca.fabs(curvature)/(1-55*ca.fabs(curvature))
@@ -266,17 +268,19 @@ def __init__(self, sr0, sr1, sr2, rho=1, g=9.81, withPnBrake=True) -> None:
 
         timeODE = 1/ca.sqrt(velocitySquared)
         velocityODE = 2*acceleration
+        positionODE = 1
 
         timeODE *= ds
         velocityODE *= ds
+        positionODE *= ds
 
-        fExplicit = ca.vertcat(timeODE, velocityODE)
+        fExplicit = ca.vertcat(timeODE, velocityODE, positionODE)
 
         # model
 
         self.ode = fExplicit
         self.acceleration = acceleration
-        self.accelerationFun = ca.Function('a', [x, u, gradient, curvature, tunnelFactor], [acceleration])
+        self.accelerationFun = ca.Function('a', [x, u, gradient, gradientLinearTerm, curvature, tunnelFactor], [acceleration])
         self.rollingResistance = rollingResistance
         self.parameters = p
         self.controls = u
@@ -356,7 +360,7 @@ def __init__(self, model, solver, optsDict={}) -> None:
             self.eval = ca.Function('xNxt', [model.states, ca.vertcat(model.controls, model.parameters), ca.MX.sym('ds')], [eval])
 
 
-    def solve(self, time, velocitySquared, ds, traction=0, pnBrake=0, gradient=0, curvature=0, tunnelFactor=0):
+    def solve(self, time, velocitySquared, ds, position=0, traction=0, pnBrake=0, gradient=0, gradientLinearTerm=0, curvature=0, tunnelFactor=0):
 
         withPnBrake = self.model.withPnBrake
 
@@ -364,9 +368,9 @@ def solve(self, time, velocitySquared, ds, traction=0, pnBrake=0, gradient=0, cu
 
             raise ValueError("Cannot define value for pneumatic braking when this brake is deactivated!")
 
-        x0 = ca.vertcat(time, velocitySquared)
+        x0 = ca.vertcat(time, velocitySquared, position)
         u0 = ca.vertcat(traction, pnBrake if withPnBrake else [])
-        p0 = ca.vertcat(gradient, curvature, tunnelFactor, ds)
+        p0 = ca.vertcat(gradient, gradientLinearTerm, curvature, tunnelFactor, ds)
         x1 = self.eval(x0, ca.vertcat(u0, p0), 1)
 
         out = {}
@@ -391,7 +395,7 @@ def initLosses(self, lossesTrFun, lossesRgbFun, totalMass, solver='CVODES'):
 
         dt = ca.MX.sym('dt')
         x = ca.vertcat(vel, ca.MX.sym('eTr'), ca.MX.sym('eBr'))
-        p = ca.vertcat(mdl.controls, mdl.parameters[0], mdl.parameters[1], dt)
+        p = ca.vertcat(mdl.controls, mdl.parameters[0], mdl.parameters[2], dt)
         xdot = ca.vertcat(velDot, energyTrDot, energyBrDot)
 
         if solver == 'RK':
diff --git a/utils.py b/utils.py
index 749c5a5..0bdbf34 100644
--- a/utils.py
+++ b/utils.py
@@ -542,6 +542,30 @@ def plotSpeedLimits(track, pos_adj, v_adj):
     plt.show()
 
 
+def plotGradients(track, pos_adj, g_adj, g_linear):
+
+    x = track.gradients.index.to_numpy(dtype=float)
+    g = track.gradients["Gradient [permil]"].to_numpy(dtype=float)
+
+    fig, ax = plt.subplots(figsize=(16, 8))
+
+    ax.step(x/1000, g, where='post', label="piecewise constant gradients")
+    ax.plot(pos_adj / 1000, g_adj, linestyle="--", label="train length averaged gradients")
+
+    g_computed = np.r_[g_adj[0], g_adj[:-1] + g_linear[:-1] * (pos_adj[1:] - pos_adj[:-1])]
+    ax.scatter(pos_adj / 1000, g_computed, marker="o", label="computed gradients")
+
+    ax.set_title("Gradients")
+    ax.set_xlabel("Position [km]")
+    ax.set_ylabel("Gradient [‰]")
+    ax.grid(True, which="both", linestyle="--", alpha=0.5)
+    ax.legend(loc="upper right")
+    ax.set_xlim(0, track.length / 1000)
+    ax.figure.tight_layout()
+
+    plt.show()
+
+
 if __name__ == '__main__':
 
     pass

From 9b6bc8cfb97e5d08d6995158f1ed83d6a61c2750 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Thu, 21 May 2026 16:06:31 +0200
Subject: [PATCH 04/31] wip 2

---
 simulations/sim_launcher.py | 2 +-
 train.py                    | 7 ++++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/simulations/sim_launcher.py b/simulations/sim_launcher.py
index 832e640..7ec173b 100644
--- a/simulations/sim_launcher.py
+++ b/simulations/sim_launcher.py
@@ -17,7 +17,7 @@
     track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
     track.updateTrainLengthDependentValues(train)
 
-    opts = {'numIntervals':400, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':1}, 'energyOptimal':True}
+    opts = {'numIntervals':400, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':0}, 'energyOptimal':True}
 
     solver = casadiSolver(train, track, opts)
 
diff --git a/train.py b/train.py
index ec7c417..ef401bc 100644
--- a/train.py
+++ b/train.py
@@ -353,7 +353,7 @@ def __init__(self, model, solver, optsDict={}) -> None:
 
                 vCurr = ca.sqrt(bf[idx])
                 vNext = ca.sqrt(bf[idx+1])
-                tApprox += 2*model.parameters[2]*(evalPoints[idx+1]-evalPoints[idx])/(vCurr + vNext)
+                tApprox += 2*model.parameters[4]*(evalPoints[idx+1]-evalPoints[idx])/(vCurr + vNext)
 
             eval = ca.vertcat(tApprox, bf[-1])
 
@@ -376,6 +376,7 @@ def solve(self, time, velocitySquared, ds, position=0, traction=0, pnBrake=0, gr
         out = {}
         out['time'] = x1[0]
         out['velSquared'] = x1[1]
+        out['position'] = x1[2]
 
         return out
 
@@ -434,7 +435,7 @@ def initRollingResistance(self, solver='CVODES'):
         mdl = self.model
 
         bDot = mdl.ode[1]
-        eDot = mdl.rollingResistance*mdl.parameters[2]
+        eDot = mdl.rollingResistance*mdl.parameters[-1]
 
         x = ca.vertcat(mdl.states[1], ca.MX.sym('e'))
         p = ca.vertcat(mdl.controls, mdl.parameters)
@@ -564,6 +565,6 @@ def checkValues(self):
     trainSpecs = Train(config={'id':'NL_intercity_VIRM6'})
     integrator = TrainIntegrator(trainSpecs.exportModel(), 'RK', optsDict={'numApproxSteps':2})
 
-    solution = integrator.solve(t0, v0**2, ds, f0, gradient=gd, curvature=cr)
+    solution = integrator.solve(t0, v0**2, ds, traction=f0, gradient=gd, curvature=cr)
 
     print(solution)

From 518780b3295eff601138d4fe0b05b8675b2924f3 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Thu, 21 May 2026 16:22:59 +0200
Subject: [PATCH 05/31] fix tunnel parameter

---
 train.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/train.py b/train.py
index 43f1a7b..402076a 100644
--- a/train.py
+++ b/train.py
@@ -347,7 +347,7 @@ def __init__(self, model, solver, optsDict={}) -> None:
 
                 vCurr = ca.sqrt(bf[idx])
                 vNext = ca.sqrt(bf[idx+1])
-                tApprox += 2*model.parameters[2]*(evalPoints[idx+1]-evalPoints[idx])/(vCurr + vNext)
+                tApprox += 2*model.parameters[-1]*(evalPoints[idx+1]-evalPoints[idx])/(vCurr + vNext)
 
             eval = ca.vertcat(tApprox, bf[-1])
 
@@ -389,7 +389,7 @@ def initLosses(self, lossesTrFun, lossesRgbFun, totalMass, solver='CVODES'):
 
         dt = ca.MX.sym('dt')
         x = ca.vertcat(vel, ca.MX.sym('eTr'), ca.MX.sym('eBr'))
-        p = ca.vertcat(mdl.controls, mdl.parameters[0], mdl.parameters[1], dt)
+        p = ca.vertcat(mdl.controls, mdl.parameters[0], mdl.parameters[1], mdl.parameters[2], dt)
         xdot = ca.vertcat(velDot, energyTrDot, energyBrDot)
 
         if solver == 'RK':
@@ -428,7 +428,7 @@ def initRollingResistance(self, solver='CVODES'):
         mdl = self.model
 
         bDot = mdl.ode[1]
-        eDot = mdl.rollingResistance*mdl.parameters[2]
+        eDot = mdl.rollingResistance*mdl.parameters[-1]
 
         x = ca.vertcat(mdl.states[1], ca.MX.sym('e'))
         p = ca.vertcat(mdl.controls, mdl.parameters)

From 3c32acd5889107352113d1b64d55bd3907f6230a Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Thu, 21 May 2026 16:52:12 +0200
Subject: [PATCH 06/31] gradients length dependent

---
 ocp.py                      |  5 +++--
 simulations/sim_launcher.py |  2 +-
 track.py                    |  2 +-
 train.py                    | 16 ++++++++--------
 4 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/ocp.py b/ocp.py
index 658f264..f91b14b 100644
--- a/ocp.py
+++ b/ocp.py
@@ -193,13 +193,14 @@ def __init__(self, train, track, optsDict={}):
 
                 # gradient and curvature of current index
                 grad = self.points.iloc[i]['Gradient [permil]']/1e3
-                gradLinearTerm = self.points.iloc[i]["Gradient linear term [permil]"]/1e3
+                gradLinearTerm = 0
+                # gradLinearTerm = self.points.iloc[i]["Gradient linear term [permil]"]/1e3
                 curv = self.points.iloc[i]['Curvature [1/m]']
                 crossSection = self.points.iloc[i]['CrossSection [m^2]']
                 tunnelFactor = computeTunnelFactor(crossSection, train)
 
                 # acceleration constraints
-                g += [trainModel.accelerationFun(ca.vertcat(time[i], velSq[i]), ca.vcat(u), grad, gradLinearTerm, curv, tunnelFactor)]
+                g += [trainModel.accelerationFun(ca.vertcat(time[i], velSq[i], 0), ca.vcat(u), grad, gradLinearTerm, curv, tunnelFactor)]
                 lbg += [accMin]
                 ubg += [accMax]
 
diff --git a/simulations/sim_launcher.py b/simulations/sim_launcher.py
index 7ec173b..97e8366 100644
--- a/simulations/sim_launcher.py
+++ b/simulations/sim_launcher.py
@@ -17,7 +17,7 @@
     track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
     track.updateTrainLengthDependentValues(train)
 
-    opts = {'numIntervals':400, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':0}, 'energyOptimal':True}
+    opts = {'numIntervals':600, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':0}, 'energyOptimal':True}
 
     solver = casadiSolver(train, track, opts)
 
diff --git a/track.py b/track.py
index 3387420..afd4e8c 100644
--- a/track.py
+++ b/track.py
@@ -513,7 +513,7 @@ def crop(dfIn):
     def updateTrainLengthDependentValues(self, train):
 
         self.updateSpeedLimitsToTrainLength(train.length)
-        self.updateGradientsToTrainLength(train.length)
+        # self.updateGradientsToTrainLength(train.length)
 
 
     def updateSpeedLimitsToTrainLength(self, trainLength):
diff --git a/train.py b/train.py
index 7e32ec3..3091447 100644
--- a/train.py
+++ b/train.py
@@ -247,7 +247,7 @@ def __init__(self, sr0, sr1, sr2, rho=1, g=9.81, withPnBrake=True) -> None:
         # parameters
 
         gradient = ca.MX.sym('gradient')                        # [-]     -> values between [0,0.2]
-        gradientLinearTerm = ca.MX.sym('gradientLinearTerm')    # [-]
+        gradientLinearTerm = ca.MX.sym('gradientLinearTerm')    # [1/m]
         curvature = ca.MX.sym('curvature')                      # [1/m]   -> values between [0,0.004]
         tunnelFactor = ca.MX.sym('tunnelFactor')                # [1/m]
         ds = ca.MX.sym('ds')
@@ -311,7 +311,7 @@ def __init__(self, model, solver, optsDict={}) -> None:
 
             opts = OptionsRK(optsDict)
 
-            ode = model.ode if opts.numApproxSteps == 0 else model.ode[1]
+            ode = model.ode if opts.numApproxSteps == 0 else model.ode[1] # todo: fix
             states = model.states if opts.numApproxSteps == 0 else model.states[1]
 
             self.eval = ca.simpleRK(ca.Function('ode', [states, params], [ode]), opts.numSteps, opts.order)
@@ -396,7 +396,7 @@ def initLosses(self, lossesTrFun, lossesRgbFun, totalMass, solver='CVODES'):
 
         dt = ca.MX.sym('dt')
         x = ca.vertcat(vel, ca.MX.sym('eTr'), ca.MX.sym('eBr'))
-        p = ca.vertcat(mdl.controls, mdl.parameters[0], mdl.parameters[1], mdl.parameters[2], dt)
+        p = ca.vertcat(mdl.controls, mdl.parameters[0], mdl.parameters[1], mdl.parameters[2], mdl.parameters[3], dt)
         xdot = ca.vertcat(velDot, energyTrDot, energyBrDot)
 
         if solver == 'RK':
@@ -419,11 +419,11 @@ def initLosses(self, lossesTrFun, lossesRgbFun, totalMass, solver='CVODES'):
             raise ValueError("Unknown solver!")
 
 
-    def calcLosses(self, velocity, dt, traction=0, pnBrake=0, gradient=0, curvature=0, tunnelFactor=0):
+    def calcLosses(self, velocity, dt, traction=0, pnBrake=0, gradient=0, gradientLinearTerm=0, curvature=0, tunnelFactor=0):
 
         mdl = self.model
 
-        out = self.lossesIntegrator(ca.vertcat(velocity, 0, 0), ca.vertcat(traction, pnBrake if mdl.withPnBrake else [], gradient, curvature, tunnelFactor), dt)
+        out = self.lossesIntegrator(ca.vertcat(velocity, 0, 0), ca.vertcat(traction, pnBrake if mdl.withPnBrake else [], gradient, gradientLinearTerm, curvature, tunnelFactor), dt)
 
         lossesTr, lossesRgb = out[1], out[2]
 
@@ -459,12 +459,12 @@ def initRollingResistance(self, solver='CVODES'):
             raise ValueError("Unknown solver!")
 
 
-    def calcRollingResistance(self, velocity, ds, traction=0, pnBrake=0, gradient=0, curvature=0, tunnelFactor=0):
+    def calcRollingResistance(self, velocity, ds, traction=0, pnBrake=0, gradient=0, gradientLinearTerm=0, curvature=0, tunnelFactor=0):
 
         mdl = self.model
 
-        out = self.rollingResistanceIntegrator(ca.vertcat(velocity**2, 0), ca.vertcat(traction, pnBrake if mdl.withPnBrake else [], gradient,
-                                                                                      curvature, tunnelFactor, ds), 1)
+        out = self.rollingResistanceIntegrator(ca.vertcat(velocity**2, 0), ca.vertcat(traction, pnBrake if mdl.withPnBrake else [],
+                                                                                      gradient, gradientLinearTerm, curvature, tunnelFactor, ds), 1)
 
         losses = out[1]
 

From 9776a5892b998b96eb0ade128eb54a793e2ecbf3 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Fri, 22 May 2026 08:45:38 +0200
Subject: [PATCH 07/31] fixed time approx

---
 ocp.py   |  7 ++++---
 track.py |  6 +++---
 train.py | 34 +++++++++++++++++++++-------------
 3 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/ocp.py b/ocp.py
index f91b14b..2a30f6e 100644
--- a/ocp.py
+++ b/ocp.py
@@ -193,8 +193,8 @@ def __init__(self, train, track, optsDict={}):
 
                 # gradient and curvature of current index
                 grad = self.points.iloc[i]['Gradient [permil]']/1e3
-                gradLinearTerm = 0
-                # gradLinearTerm = self.points.iloc[i]["Gradient linear term [permil]"]/1e3
+                # gradLinearTerm = 0
+                gradLinearTerm = self.points.iloc[i]["Gradient linear term [permil]"]/1e3
                 curv = self.points.iloc[i]['Curvature [1/m]']
                 crossSection = self.points.iloc[i]['CrossSection [m^2]']
                 tunnelFactor = computeTunnelFactor(crossSection, train)
@@ -326,7 +326,8 @@ def solve(self, terminalTime, initialTime=0, terminalVelocity=1, initialVelocity
         # initial guess
         # NOTE: good idea vel0 to be compatible with f0 (power-wise) to avoid nans at first iteration
 
-        vel0 = (60/3.6)**2
+        vel_avg = (self.points.index[-1]-self.points.index[0])/(terminalTime-initialTime)
+        vel0 = vel_avg*vel_avg
         dt = (terminalTime - initialTime)/self.numIntervals
         t0 = initialTime
 
diff --git a/track.py b/track.py
index afd4e8c..af25450 100644
--- a/track.py
+++ b/track.py
@@ -513,7 +513,7 @@ def crop(dfIn):
     def updateTrainLengthDependentValues(self, train):
 
         self.updateSpeedLimitsToTrainLength(train.length)
-        # self.updateGradientsToTrainLength(train.length)
+        self.updateGradientsToTrainLength(train.length)
 
 
     def updateSpeedLimitsToTrainLength(self, trainLength):
@@ -545,7 +545,7 @@ def updateSpeedLimitsToTrainLength(self, trainLength):
                 pos_adj.append(new_pos)
                 v_adj.append(v[i])
 
-            plotSpeedLimits(self, np.asarray(pos_adj, dtype=float), np.asarray(v_adj, dtype=float))
+            # plotSpeedLimits(self, np.asarray(pos_adj, dtype=float), np.asarray(v_adj, dtype=float))
 
             self.speedLimits = pd.DataFrame(
                 {"Speed limit [m/s]": v_adj},
@@ -588,7 +588,7 @@ def updateGradientsToTrainLength(self, trainLength):
                 g_adj.append(currentGradient)
                 g_linear.append(currentLinearTerm)
 
-            plotGradients(self, np.asarray(pos_adj, dtype=float), np.asarray(g_adj, dtype=float), np.asarray(g_linear, dtype=float))
+            # plotGradients(self, np.asarray(pos_adj, dtype=float), np.asarray(g_adj, dtype=float), np.asarray(g_linear, dtype=float))
 
             self.gradients = pd.DataFrame(
                 {"Gradient [permil]": g_adj, "Gradient linear term [permil]": g_linear},
diff --git a/train.py b/train.py
index 3091447..37e2457 100644
--- a/train.py
+++ b/train.py
@@ -258,7 +258,7 @@ def __init__(self, sr0, sr1, sr2, rho=1, g=9.81, withPnBrake=True) -> None:
 
         rollingResistance = sr0 + sr1*ca.sqrt(velocitySquared) + sr2*velocitySquared                # [N/kg]
         gradientResistance = g*(1/rho)*(gradient+gradientLinearTerm*position)                                                  # [N/kg]
-        curvatureResistance = ca.if_else(ca.fabs(curvature)<=1/300,
+        curvatureResistance = (1/rho)*ca.if_else(ca.fabs(curvature)<=1/300,
                                             g*0.5*ca.fabs(curvature)/(1-30*ca.fabs(curvature)),
                                             g*0.65*ca.fabs(curvature)/(1-55*ca.fabs(curvature))
                                          )                                                          # [N/kg]
@@ -311,8 +311,12 @@ def __init__(self, model, solver, optsDict={}) -> None:
 
             opts = OptionsRK(optsDict)
 
-            ode = model.ode if opts.numApproxSteps == 0 else model.ode[1] # todo: fix
-            states = model.states if opts.numApproxSteps == 0 else model.states[1]
+            if opts.numApproxSteps == 0:
+                ode = model.ode
+                states = model.states
+            else:
+                ode = ca.vertcat(model.ode[1], model.ode[2])
+                states = ca.vertcat(model.states[1], model.states[2])
 
             self.eval = ca.simpleRK(ca.Function('ode', [states, params], [ode]), opts.numSteps, opts.order)
 
@@ -320,10 +324,14 @@ def __init__(self, model, solver, optsDict={}) -> None:
 
             opts = OptionsIRK(optsDict)
 
-            ode = model.ode if opts.numApproxSteps == 0 else model.ode[1]
-            states = model.states if opts.numApproxSteps == 0 else model.states[1]
+            if opts.numApproxSteps == 0:
+                ode = model.ode
+                states = model.states
+            else:
+                ode = ca.vertcat(model.ode[1], model.ode[2])
+                states = ca.vertcat(model.states[1], model.states[2])
 
-            self.eval = ca.simpleIRK(ca.Function('ode', [states, params], [ode]), opts.numSteps, opts.order, opts.collMethod, 'fast_newton', {'max_iter':opts.maxIter, 'jit':opts.jit, 'error_on_fail':False})
+            self.eval = ca.simpleIRK(ca.Function('ode', [states, params], [ode]), opts.numSteps, opts.order, opts.collMethod, 'fast_newton', {'max_iter': opts.maxIter, 'jit': opts.jit, 'error_on_fail': False})
 
         elif solver == 'CVODES':
 
@@ -333,7 +341,7 @@ def __init__(self, model, solver, optsDict={}) -> None:
             states = model.states
 
             t0, tf = 0, 1
-            cvodesFun = ca.integrator('integrator', 'cvodes', {'x':model.states, 'p':params, 'ode':model.ode}, t0, tf, {'abstol':opts.absTol, 'reltol':opts.relTol})
+            cvodesFun = ca.integrator('integrator', 'cvodes', {'x': model.states, 'p': params, 'ode': model.ode}, t0, tf, {'abstol': opts.absTol, 'reltol': opts.relTol})
 
             self.eval = lambda x, p, dummy: cvodesFun(x0=x, p=p)['xf']
 
@@ -343,19 +351,19 @@ def __init__(self, model, solver, optsDict={}) -> None:
 
             evalPoints = [0] + [i/ns for i in range(1, ns+1)]
 
-            b0 = model.states[1]
+            z0 = ca.vertcat(model.states[1], model.states[2])
             p0 = ca.vertcat(model.controls, model.parameters)
-            bf = self.eval(b0, p0, ca.hcat(evalPoints))
+            zf = self.eval(z0, p0, ca.hcat(evalPoints))
 
             tApprox = model.states[0]
 
+            epsVelSq = 0.0001
             for idx in range(ns):
-
-                vCurr = ca.sqrt(bf[idx])
-                vNext = ca.sqrt(bf[idx+1])
+                vCurr = ca.sqrt(ca.fmax(zf[0, idx], epsVelSq))
+                vNext = ca.sqrt(ca.fmax(zf[0, idx + 1], epsVelSq))
                 tApprox += 2*model.parameters[-1]*(evalPoints[idx+1]-evalPoints[idx])/(vCurr + vNext)
 
-            eval = ca.vertcat(tApprox, bf[-1])
+            eval = ca.vertcat(tApprox, zf[0, ns], zf[1, ns])
 
             self.eval = ca.Function('xNxt', [model.states, ca.vertcat(model.controls, model.parameters), ca.MX.sym('ds')], [eval])
 

From 9b9c212bddddbd220878b803e9d1b544b914ef39 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Fri, 22 May 2026 09:13:53 +0200
Subject: [PATCH 08/31] curvature train length dependent

---
 ocp.py                      |  3 ++-
 simulations/sim_launcher.py |  2 +-
 track.py                    | 44 ++++++++++++++++++++++++++++++++++++-
 train.py                    | 22 +++++++++----------
 utils.py                    | 24 ++++++++++++++++++++
 5 files changed, 80 insertions(+), 15 deletions(-)

diff --git a/ocp.py b/ocp.py
index 2a30f6e..00c9495 100644
--- a/ocp.py
+++ b/ocp.py
@@ -196,11 +196,12 @@ def __init__(self, train, track, optsDict={}):
                 # gradLinearTerm = 0
                 gradLinearTerm = self.points.iloc[i]["Gradient linear term [permil]"]/1e3
                 curv = self.points.iloc[i]['Curvature [1/m]']
+                curvLinearTerm = self.points.iloc[i]["Curvature linear term [1/m^2]"]
                 crossSection = self.points.iloc[i]['CrossSection [m^2]']
                 tunnelFactor = computeTunnelFactor(crossSection, train)
 
                 # acceleration constraints
-                g += [trainModel.accelerationFun(ca.vertcat(time[i], velSq[i], 0), ca.vcat(u), grad, gradLinearTerm, curv, tunnelFactor)]
+                g += [trainModel.accelerationFun(ca.vertcat(time[i], velSq[i], 0), ca.vcat(u), grad, gradLinearTerm, curv, curvLinearTerm, tunnelFactor)]
                 lbg += [accMin]
                 ubg += [accMax]
 
diff --git a/simulations/sim_launcher.py b/simulations/sim_launcher.py
index 97e8366..55eb1c0 100644
--- a/simulations/sim_launcher.py
+++ b/simulations/sim_launcher.py
@@ -17,7 +17,7 @@
     track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
     track.updateTrainLengthDependentValues(train)
 
-    opts = {'numIntervals':600, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':0}, 'energyOptimal':True}
+    opts = {'numIntervals':600, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':1}, 'energyOptimal':True}
 
     solver = casadiSolver(train, track, opts)
 
diff --git a/track.py b/track.py
index af25450..e8e6890 100644
--- a/track.py
+++ b/track.py
@@ -5,7 +5,8 @@
 import numpy as np
 import pandas as pd
 
-from utils import checkTTOBenchVersion, convertUnit, pickEquallySpacedPoints, plotSpeedLimits, plotGradients
+from utils import checkTTOBenchVersion, convertUnit, pickEquallySpacedPoints, plotSpeedLimits, plotGradients, \
+    plotCurvatures
 
 
 def importTuples(tuples, xLabel, yLabels):
@@ -514,6 +515,7 @@ def updateTrainLengthDependentValues(self, train):
 
         self.updateSpeedLimitsToTrainLength(train.length)
         self.updateGradientsToTrainLength(train.length)
+        self.updateCurvaturesToTrainLength(train.length)
 
 
     def updateSpeedLimitsToTrainLength(self, trainLength):
@@ -595,7 +597,47 @@ def updateGradientsToTrainLength(self, trainLength):
                 index=pos_adj,
             )
 
+    def updateCurvaturesToTrainLength(self, trainLength):
 
+        c = self.curvatures["Curvature [1/m]"].to_numpy(dtype=float)
+        pos = self.curvatures.index.to_numpy(dtype=float)
+        slopes = np.r_[0, (c[1:] - c[:-1]) / trainLength]
+
+        if len(pos) > 1:
+
+            pos_adj = np.sort(np.r_[pos, pos + trainLength])
+
+            pos_adj = pos_adj[pos_adj < self.length]
+
+            c_adj = [c[0]]
+            c_linear = [0]
+
+            for idx in range(1,len(pos_adj)):
+
+                currentPosition = pos_adj[idx]
+                previousPosition = pos_adj[idx - 1]
+
+                currentCurvature = c_adj[idx-1] + (currentPosition-previousPosition)*c_linear[idx-1]
+
+                epsilon = 0.001
+                list_indices = []
+
+                for idx2 in range(len(pos)-1):
+
+                    if currentPosition - trainLength + epsilon < pos[idx2] < currentPosition + epsilon:
+                        list_indices.append(idx2)
+
+                currentLinearTerm = sum(slopes[list_indices])
+
+                c_adj.append(currentCurvature)
+                c_linear.append(currentLinearTerm)
+
+            # plotCurvatures(self, np.asarray(pos_adj, dtype=float), np.asarray(c_adj, dtype=float), np.asarray(c_linear, dtype=float))
+
+            self.curvatures = pd.DataFrame(
+                {"Curvature [1/m]": c_adj, "Curvature linear term [1/m^2]": c_linear},
+                index=pos_adj,
+            )
 
 
 if __name__ == '__main__':
diff --git a/train.py b/train.py
index 37e2457..5a3aadc 100644
--- a/train.py
+++ b/train.py
@@ -249,19 +249,17 @@ def __init__(self, sr0, sr1, sr2, rho=1, g=9.81, withPnBrake=True) -> None:
         gradient = ca.MX.sym('gradient')                        # [-]     -> values between [0,0.2]
         gradientLinearTerm = ca.MX.sym('gradientLinearTerm')    # [1/m]
         curvature = ca.MX.sym('curvature')                      # [1/m]   -> values between [0,0.004]
+        curvatureLinearTerm = ca.MX.sym('curvatureLinearTerm')            # [1/m^2]
         tunnelFactor = ca.MX.sym('tunnelFactor')                # [1/m]
         ds = ca.MX.sym('ds')
 
-        p = ca.vertcat(gradient, gradientLinearTerm, curvature, tunnelFactor, ds)
+        p = ca.vertcat(gradient, gradientLinearTerm, curvature, curvatureLinearTerm, tunnelFactor, ds)
 
         # ODE
 
         rollingResistance = sr0 + sr1*ca.sqrt(velocitySquared) + sr2*velocitySquared                # [N/kg]
         gradientResistance = g*(1/rho)*(gradient+gradientLinearTerm*position)                                                  # [N/kg]
-        curvatureResistance = (1/rho)*ca.if_else(ca.fabs(curvature)<=1/300,
-                                            g*0.5*ca.fabs(curvature)/(1-30*ca.fabs(curvature)),
-                                            g*0.65*ca.fabs(curvature)/(1-55*ca.fabs(curvature))
-                                         )                                                          # [N/kg]
+        curvatureResistance = (1/rho)*(5.07 * (curvature+curvatureLinearTerm*position))                                                          # [N/kg]
         tunnelResistance = tunnelFactor * velocitySquared                                           # [N/kg]
 
         acceleration = traction + (pnBrake if withPnBrake else 0) - rollingResistance - gradientResistance - curvatureResistance - tunnelResistance # [m/s^2]
@@ -280,7 +278,7 @@ def __init__(self, sr0, sr1, sr2, rho=1, g=9.81, withPnBrake=True) -> None:
 
         self.ode = fExplicit
         self.acceleration = acceleration
-        self.accelerationFun = ca.Function('a', [x, u, gradient, gradientLinearTerm, curvature, tunnelFactor], [acceleration])
+        self.accelerationFun = ca.Function('a', [x, u, gradient, gradientLinearTerm, curvature, curvatureLinearTerm, tunnelFactor], [acceleration])
         self.rollingResistance = rollingResistance
         self.parameters = p
         self.controls = u
@@ -368,7 +366,7 @@ def __init__(self, model, solver, optsDict={}) -> None:
             self.eval = ca.Function('xNxt', [model.states, ca.vertcat(model.controls, model.parameters), ca.MX.sym('ds')], [eval])
 
 
-    def solve(self, time, velocitySquared, ds, position=0, traction=0, pnBrake=0, gradient=0, gradientLinearTerm=0, curvature=0, tunnelFactor=0):
+    def solve(self, time, velocitySquared, ds, position=0, traction=0, pnBrake=0, gradient=0, gradientLinearTerm=0, curvature=0, curvatureLinearTerm=0, tunnelFactor=0):
 
         withPnBrake = self.model.withPnBrake
 
@@ -378,7 +376,7 @@ def solve(self, time, velocitySquared, ds, position=0, traction=0, pnBrake=0, gr
 
         x0 = ca.vertcat(time, velocitySquared, position)
         u0 = ca.vertcat(traction, pnBrake if withPnBrake else [])
-        p0 = ca.vertcat(gradient, gradientLinearTerm, curvature, tunnelFactor, ds)
+        p0 = ca.vertcat(gradient, gradientLinearTerm, curvature, curvatureLinearTerm, tunnelFactor, ds)
         x1 = self.eval(x0, ca.vertcat(u0, p0), 1)
 
         out = {}
@@ -427,11 +425,11 @@ def initLosses(self, lossesTrFun, lossesRgbFun, totalMass, solver='CVODES'):
             raise ValueError("Unknown solver!")
 
 
-    def calcLosses(self, velocity, dt, traction=0, pnBrake=0, gradient=0, gradientLinearTerm=0, curvature=0, tunnelFactor=0):
+    def calcLosses(self, velocity, dt, traction=0, pnBrake=0, gradient=0, gradientLinearTerm=0, curvature=0, curvatureLinearTerm=0, tunnelFactor=0):
 
         mdl = self.model
 
-        out = self.lossesIntegrator(ca.vertcat(velocity, 0, 0), ca.vertcat(traction, pnBrake if mdl.withPnBrake else [], gradient, gradientLinearTerm, curvature, tunnelFactor), dt)
+        out = self.lossesIntegrator(ca.vertcat(velocity, 0, 0), ca.vertcat(traction, pnBrake if mdl.withPnBrake else [], gradient, gradientLinearTerm, curvature, curvatureLinearTerm, tunnelFactor), dt)
 
         lossesTr, lossesRgb = out[1], out[2]
 
@@ -467,12 +465,12 @@ def initRollingResistance(self, solver='CVODES'):
             raise ValueError("Unknown solver!")
 
 
-    def calcRollingResistance(self, velocity, ds, traction=0, pnBrake=0, gradient=0, gradientLinearTerm=0, curvature=0, tunnelFactor=0):
+    def calcRollingResistance(self, velocity, ds, traction=0, pnBrake=0, gradient=0, gradientLinearTerm=0, curvature=0, curvatureLinearTerm=0, tunnelFactor=0):
 
         mdl = self.model
 
         out = self.rollingResistanceIntegrator(ca.vertcat(velocity**2, 0), ca.vertcat(traction, pnBrake if mdl.withPnBrake else [],
-                                                                                      gradient, gradientLinearTerm, curvature, tunnelFactor, ds), 1)
+                                                                                      gradient, gradientLinearTerm, curvature, curvatureLinearTerm, tunnelFactor, ds), 1)
 
         losses = out[1]
 
diff --git a/utils.py b/utils.py
index 0bdbf34..92d297b 100644
--- a/utils.py
+++ b/utils.py
@@ -566,6 +566,30 @@ def plotGradients(track, pos_adj, g_adj, g_linear):
     plt.show()
 
 
+def plotCurvatures(track, pos_adj, c_adj, c_linear):
+
+    x = track.curvatures.index.to_numpy(dtype=float)
+    c = track.curvatures["Curvature [1/m]"].to_numpy(dtype=float)
+
+    fig, ax = plt.subplots(figsize=(16, 8))
+
+    ax.step(x/1000, c, where='post', label="piecewise constant curvatures")
+    ax.plot(pos_adj / 1000, c_adj, linestyle="--", label="train length averaged curvatures")
+
+    c_computed = np.r_[c_adj[0], c_adj[:-1] + c_linear[:-1] * (pos_adj[1:] - pos_adj[:-1])]
+    ax.scatter(pos_adj / 1000, c_computed, marker="o", label="computed curvatures")
+
+    ax.set_title("Curvatures")
+    ax.set_xlabel("Position [km]")
+    ax.set_ylabel("Curvature [1/m]")
+    ax.grid(True, which="both", linestyle="--", alpha=0.5)
+    ax.legend(loc="upper right")
+    ax.set_xlim(0, track.length / 1000)
+    ax.figure.tight_layout()
+
+    plt.show()
+
+
 if __name__ == '__main__':
 
     pass

From 43b6e16f2d46842a113bea1cfead8d94139e5609 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Fri, 22 May 2026 09:47:41 +0200
Subject: [PATCH 09/31] bug fixes

---
 simulations/sim_launcher.py |  2 +-
 train.py                    | 23 +++++++++++++----------
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/simulations/sim_launcher.py b/simulations/sim_launcher.py
index 55eb1c0..97e8366 100644
--- a/simulations/sim_launcher.py
+++ b/simulations/sim_launcher.py
@@ -17,7 +17,7 @@
     track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
     track.updateTrainLengthDependentValues(train)
 
-    opts = {'numIntervals':600, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':1}, 'energyOptimal':True}
+    opts = {'numIntervals':600, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':0}, 'energyOptimal':True}
 
     solver = casadiSolver(train, track, opts)
 
diff --git a/train.py b/train.py
index 5a3aadc..a43409d 100644
--- a/train.py
+++ b/train.py
@@ -394,16 +394,18 @@ def initLosses(self, lossesTrFun, lossesRgbFun, totalMass, solver='CVODES'):
 
         mdl = self.model
         vel = ca.MX.sym('v')
+        pos = ca.MX.sym('pos')
 
         velDot = ca.substitute(mdl.acceleration, mdl.states[1], vel**2)
+        velDot = ca.substitute(velDot, mdl.states[2], pos)
 
         energyTrDot = lossesTrFun(mdl.controls[0]*totalMass, vel)/totalMass  # tractive energy
         energyBrDot = lossesRgbFun(mdl.controls[0]*totalMass, vel)/totalMass  # braking energy
 
         dt = ca.MX.sym('dt')
-        x = ca.vertcat(vel, ca.MX.sym('eTr'), ca.MX.sym('eBr'))
-        p = ca.vertcat(mdl.controls, mdl.parameters[0], mdl.parameters[1], mdl.parameters[2], mdl.parameters[3], dt)
-        xdot = ca.vertcat(velDot, energyTrDot, energyBrDot)
+        x = ca.vertcat(vel, pos, ca.MX.sym('eTr'), ca.MX.sym('eBr'))
+        p = ca.vertcat(mdl.controls, mdl.parameters[0], mdl.parameters[1], mdl.parameters[2], mdl.parameters[3], mdl.parameters[4], dt)
+        xdot = ca.vertcat(velDot, vel, energyTrDot, energyBrDot)
 
         if solver == 'RK':
 
@@ -425,11 +427,11 @@ def initLosses(self, lossesTrFun, lossesRgbFun, totalMass, solver='CVODES'):
             raise ValueError("Unknown solver!")
 
 
-    def calcLosses(self, velocity, dt, traction=0, pnBrake=0, gradient=0, gradientLinearTerm=0, curvature=0, curvatureLinearTerm=0, tunnelFactor=0):
+    def calcLosses(self, velocity, dt, position=0, traction=0, pnBrake=0, gradient=0, gradientLinearTerm=0, curvature=0, curvatureLinearTerm=0, tunnelFactor=0):
 
         mdl = self.model
 
-        out = self.lossesIntegrator(ca.vertcat(velocity, 0, 0), ca.vertcat(traction, pnBrake if mdl.withPnBrake else [], gradient, gradientLinearTerm, curvature, curvatureLinearTerm, tunnelFactor), dt)
+        out = self.lossesIntegrator(ca.vertcat(velocity, position, 0, 0), ca.vertcat(traction, pnBrake if mdl.withPnBrake else [], gradient, gradientLinearTerm, curvature, curvatureLinearTerm, tunnelFactor), dt)
 
         lossesTr, lossesRgb = out[1], out[2]
 
@@ -441,11 +443,12 @@ def initRollingResistance(self, solver='CVODES'):
         mdl = self.model
 
         bDot = mdl.ode[1]
+        sDot = mdl.ode[2]
         eDot = mdl.rollingResistance*mdl.parameters[-1]
 
-        x = ca.vertcat(mdl.states[1], ca.MX.sym('e'))
+        x = ca.vertcat(mdl.states[1], mdl.states[2], ca.MX.sym('e'))
         p = ca.vertcat(mdl.controls, mdl.parameters)
-        xdot = ca.vertcat(bDot, eDot)
+        xdot = ca.vertcat(bDot, sDot, eDot)
 
         if solver == 'RK':
 
@@ -465,14 +468,14 @@ def initRollingResistance(self, solver='CVODES'):
             raise ValueError("Unknown solver!")
 
 
-    def calcRollingResistance(self, velocity, ds, traction=0, pnBrake=0, gradient=0, gradientLinearTerm=0, curvature=0, curvatureLinearTerm=0, tunnelFactor=0):
+    def calcRollingResistance(self, velocity, ds, position=0, traction=0, pnBrake=0, gradient=0, gradientLinearTerm=0, curvature=0, curvatureLinearTerm=0, tunnelFactor=0):
 
         mdl = self.model
 
-        out = self.rollingResistanceIntegrator(ca.vertcat(velocity**2, 0), ca.vertcat(traction, pnBrake if mdl.withPnBrake else [],
+        out = self.rollingResistanceIntegrator(ca.vertcat(velocity**2, position, 0), ca.vertcat(traction, pnBrake if mdl.withPnBrake else [],
                                                                                       gradient, gradientLinearTerm, curvature, curvatureLinearTerm, tunnelFactor, ds), 1)
 
-        losses = out[1]
+        losses = out[2]
 
         return losses, ca.sqrt(out[0])
 

From 0e20613cb5605f221b2f924754d5eaaaf5a3462b Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Fri, 22 May 2026 09:57:43 +0200
Subject: [PATCH 10/31] bug fix

---
 train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/train.py b/train.py
index a43409d..ab1a4c8 100644
--- a/train.py
+++ b/train.py
@@ -433,7 +433,7 @@ def calcLosses(self, velocity, dt, position=0, traction=0, pnBrake=0, gradient=0
 
         out = self.lossesIntegrator(ca.vertcat(velocity, position, 0, 0), ca.vertcat(traction, pnBrake if mdl.withPnBrake else [], gradient, gradientLinearTerm, curvature, curvatureLinearTerm, tunnelFactor), dt)
 
-        lossesTr, lossesRgb = out[1], out[2]
+        lossesTr, lossesRgb = out[2], out[3]
 
         return lossesTr, lossesRgb
 

From 0089887bc48889915d386cc0abab236560c8a964 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Tue, 26 May 2026 12:48:44 +0200
Subject: [PATCH 11/31] length dependent constant values adjusted

---
 track.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/track.py b/track.py
index e8e6890..3c299d9 100644
--- a/track.py
+++ b/track.py
@@ -105,6 +105,27 @@ def computeDiscretizationPoints(track, numIntervals):
 
         raise ValueError("Wrong number of computed discretization intervals!")
 
+    # adapt constant track attribute terms
+
+    positions = df3.index.to_numpy(dtype=float)
+    grads = [df3["Gradient [permil]"].iloc[0]]
+    curvs = [df3["Curvature [1/m]"].iloc[0]]
+
+    for idx in range(1, numIntervals+1):
+
+        if np.isclose(df3["Gradient [permil]"].iloc[idx-1], df3["Gradient [permil]"].iloc[idx]):
+            grads.append(grads[-1] + (positions[idx] -  positions[idx - 1]) * df3["Gradient linear term [permil]"].iloc[idx - 1])
+        else:
+            grads.append(df3["Gradient [permil]"].iloc[idx])
+
+        if np.isclose(df3["Curvature [1/m]"].iloc[idx - 1], df3["Curvature [1/m]"].iloc[idx]):
+            curvs.append(curvs[-1] + (positions[idx] - positions[idx - 1]) * df3["Curvature linear term [1/m^2]"].iloc[idx - 1])
+        else:
+            curvs.append(df3["Curvature [1/m]"].iloc[idx])
+
+    df3["Gradient [permil]"] = grads
+    df3["Curvature [1/m]"] = curvs
+
     return df3
 
 

From b1b15663ac7a0efdb97ac5238255182e28ee411c Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Tue, 26 May 2026 13:50:44 +0200
Subject: [PATCH 12/31] tunnel resistance made to list

---
 ocp.py   |  4 +++-
 train.py | 22 +++++++++++++++++-----
 utils.py | 41 ++++++++++++++++++++++++++++++++++++-----
 3 files changed, 56 insertions(+), 11 deletions(-)

diff --git a/ocp.py b/ocp.py
index 00c9495..43b80b6 100644
--- a/ocp.py
+++ b/ocp.py
@@ -27,6 +27,8 @@ def __init__(self, paramsDict):
 
         self.integrateLosses = False  # integrate losses or take mid-point rule
 
+        self.chooseClosestTunnelCrossSection = True # if exact tunnel cross section is not specified in train tunnel resistances, choose the closest value
+
         super().__init__(paramsDict)
 
 
@@ -198,7 +200,7 @@ def __init__(self, train, track, optsDict={}):
                 curv = self.points.iloc[i]['Curvature [1/m]']
                 curvLinearTerm = self.points.iloc[i]["Curvature linear term [1/m^2]"]
                 crossSection = self.points.iloc[i]['CrossSection [m^2]']
-                tunnelFactor = computeTunnelFactor(crossSection, train)
+                tunnelFactor = computeTunnelFactor(crossSection, train, opts)
 
                 # acceleration constraints
                 g += [trainModel.accelerationFun(ca.vertcat(time[i], velSq[i], 0), ca.vcat(u), grad, gradLinearTerm, curv, curvLinearTerm, tunnelFactor)]
diff --git a/train.py b/train.py
index ab1a4c8..f82e35b 100644
--- a/train.py
+++ b/train.py
@@ -30,7 +30,7 @@ def __init__(self, config, pathJSON='trains') -> None:
 
             data = json.load(file)
 
-        checkTTOBenchVersion(data, ['1.1', '1.2', '1.3'])
+        checkTTOBenchVersion(data, ['1.1', '1.2', '1.3', '1.4'])
 
         # overwrite json data with config values if applicable
 
@@ -97,10 +97,6 @@ def __init__(self, config, pathJSON='trains') -> None:
 
         self.r2 = convertUnit(data['rolling resistance r2']['value'], data['rolling resistance r2']['unit'])  # quadratic term [N/(m/s)^2]
 
-        self.t_24 = convertUnit(data['tunnel resistance 24 m^2']['value'], data['tunnel resistance 24 m^2']['unit'])  # [kg/m]
-
-        self.t_40 = convertUnit(data['tunnel resistance 40 m^2']['value'], data['tunnel resistance 40 m^2']['unit'])  # [kg/m]
-
         # TODO: unify with case of dynamic efficiency
         if 'efficiency traction' in data or 'efficiency reg brake' in data:
 
@@ -115,6 +111,22 @@ def __init__(self, config, pathJSON='trains') -> None:
             self.etaTraction = convertUnit(data['efficiency traction']['value'], data['efficiency traction']['unit'])
             self.etaRgBrake = convertUnit(data['efficiency reg brake']['value'], data['efficiency reg brake']['unit'])
 
+        if "tunnel resistance" in data:
+
+            tunnel_data = data["tunnel resistance"] # additive aerodynamic tunnel drag [kg/m] as dict per tunnel cross section [m^2]
+
+            self.tunnelCoefficients = {
+                convertUnit(item["cross section"], tunnel_data["cross section unit"]): convertUnit(
+                    item["value"],
+                    tunnel_data["unit"]
+                )
+                for item in tunnel_data["values"]
+            }
+
+        else:
+
+            self.tunnelCoefficients = {}
+
         self.checkFields()
 
 
diff --git a/utils.py b/utils.py
index 92d297b..c93cf8c 100644
--- a/utils.py
+++ b/utils.py
@@ -482,16 +482,42 @@ def latexify():
     return latexFound
 
 
-def computeTunnelFactor(cross_section, train):
+def computeTunnelFactor(cross_section, train, opts):
+
+    if cross_section == float("inf"):
 
-    if cross_section == 0:
         return 0 # no tunnel
 
     total_mass = train.mass * train.rho
-    t_24 = train.t_24
-    t_40 = train.t_40
+    tunnelCoefficients = train.tunnelCoefficients
+
+    if not tunnelCoefficients:
+        raise ValueError("Tunnel cross section was specified, but train has no tunnel resistance data.")
+
+    availableCrossSections = list(tunnelCoefficients.keys())
+
+    if opts.chooseClosestTunnelCrossSection:
+
+        closestCrossSection = min(availableCrossSections, key=lambda cs: abs(cs - cross_section))
+        tunnelCoefficient = tunnelCoefficients[closestCrossSection]
+
+    else:
+
+        if cross_section not in tunnelCoefficients:
 
-    return t_24/total_mass if cross_section < 30 else t_40/total_mass
+            raise ValueError(
+                "Tunnel resistance coefficient for cross section {} m^2 is not available. "
+                "Available cross sections are: {}. "
+                "Set chooseClosestTunnelCrossSection=True to use the closest available value.".format(
+                    cross_section,
+                    availableCrossSections
+                )
+            )
+
+        tunnelCoefficient = tunnelCoefficients[cross_section]
+
+
+    return tunnelCoefficient/total_mass
 
 
 def pickEquallySpacedPoints(startPoint, endPoint, numIntervals, requiredPoints):
@@ -499,19 +525,24 @@ def pickEquallySpacedPoints(startPoint, endPoint, numIntervals, requiredPoints):
     np.random.seed(42)
 
     if len(requiredPoints) > numIntervals + 1:
+
         raise ValueError(f"Too many required points ({len(requiredPoints)}) for N.")
 
     num_of_remaining_points = numIntervals + 1 - len(requiredPoints)
     m = 1 # number of points to oversample
 
     while True:
+
         cand = np.linspace(startPoint, endPoint, num_of_remaining_points + m + 2)[1:-1]  # oversample to avoid overlaps
         cand = np.round(cand, 0)
         out = np.unique(np.r_[requiredPoints, cand])
         cand_without_required = out[~np.isin(out, requiredPoints)]
+
         if len(cand_without_required) >= num_of_remaining_points:
+
             picked_points = np.random.choice(cand_without_required, size=num_of_remaining_points, replace=False)
             return picked_points
+
         m *= 2
 
 

From 52fdc5e11c6c02d1e25eebc42ab4721431f30248 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Tue, 26 May 2026 14:21:24 +0200
Subject: [PATCH 13/31] update tunnel resistance format

---
 simulations/sim_launcher.py                   |  2 +-
 train.py                                      | 13 ++-
 trains/Flirt_Tpf.json                         | 96 +++++++++++++++++++
 .../speedLimit.py                             |  0
 .../tunnelResistance/tunnelResistance.py      | 18 ++++
 5 files changed, 123 insertions(+), 6 deletions(-)
 create mode 100644 trains/Flirt_Tpf.json
 create mode 100644 unitTests/trainLengthDependentTrackAttributes/speedLimit.py
 create mode 100644 unitTests/tunnelResistance/tunnelResistance.py

diff --git a/simulations/sim_launcher.py b/simulations/sim_launcher.py
index 97e8366..558243c 100644
--- a/simulations/sim_launcher.py
+++ b/simulations/sim_launcher.py
@@ -10,7 +10,7 @@
     endPosition = 20000     # [m]
     duration = 60*20        # [s]
 
-    train = Train(config={'id':'SBB_Flirt_2'}, pathJSON='../trains')
+    train = Train(config={'id':'Flirt_Tpf'}, pathJSON='../trains')
 
     track = Track(config={'id':'CH_ZH_LU'}, pathJSON='../tracks')
     # track = Track(config={'id':'CH_StGallen_Wil'}, pathJSON='../tracks')
diff --git a/train.py b/train.py
index f82e35b..37b92af 100644
--- a/train.py
+++ b/train.py
@@ -113,14 +113,17 @@ def __init__(self, config, pathJSON='trains') -> None:
 
         if "tunnel resistance" in data:
 
-            tunnel_data = data["tunnel resistance"] # additive aerodynamic tunnel drag [kg/m] as dict per tunnel cross section [m^2]
+            tunnel_data = data["tunnel resistance"] # additive aerodynamic tunnel drag as dict per tunnel cross section
+
+            cross_section_unit = tunnel_data["units"]["cross section"] # tunnel cross section [m^2]
+            resistance_unit = tunnel_data["units"]["resistance (or similar)"] # resistance coefficient [kg/m]
 
             self.tunnelCoefficients = {
-                convertUnit(item["cross section"], tunnel_data["cross section unit"]): convertUnit(
-                    item["value"],
-                    tunnel_data["unit"]
+                convertUnit(cross_section, cross_section_unit): convertUnit(
+                    resistance,
+                    resistance_unit
                 )
-                for item in tunnel_data["values"]
+                for cross_section, resistance in tunnel_data["values"]
             }
 
         else:
diff --git a/trains/Flirt_Tpf.json b/trains/Flirt_Tpf.json
new file mode 100644
index 0000000..a97c321
--- /dev/null
+++ b/trains/Flirt_Tpf.json
@@ -0,0 +1,96 @@
+{
+    "metadata": {
+        "id": "CH_Stadler_FLIRT_TPF",
+        "created by": "Dimitris Kouzoupis",
+        "library version": "TTOBench v1.4",
+        "license": "BSD 2-Clause License"
+    },
+    "num seats": {
+        "unit": "-",
+        "value":167
+    },
+    "num coaches": {
+        "unit": "-",
+        "value": 4
+    },
+    "length": {
+        "unit": "m",
+        "value": 58.6
+    },
+    "mass": {
+        "unit": "kg",
+        "value": 122000.0
+    },
+    "rho": {
+        "unit": "%",
+        "value": 10
+    },
+    "max traction power": {
+        "unit": "kW",
+        "value": 2600.0
+    },
+    "max reg braking power": {
+        "unit": "kW",
+        "value": 2600.0
+    },
+    "max traction force": {
+        "unit": "kN",
+        "value": 200.0
+    },
+    "max reg braking force": {
+        "unit": "kN",
+        "value": 200.0
+    },
+    "max pn braking force": {
+        "unit": "kN",
+        "value": 1220.0
+    },
+    "max acceleration": {
+        "unit": "m/s^2",
+        "value": 1.1
+    },
+    "max deceleration": {
+        "unit": "m/s^2",
+        "value": 1.1
+    },
+    "max speed": {
+        "unit": "km/h",
+        "value": 160
+    },
+    "rolling resistance r0": {
+        "unit": "kN",
+        "value": 2.37888
+    },
+    "rolling resistance r1": {
+        "unit": "kN/(km/h)",
+        "value": 0.01698165
+    },
+    "rolling resistance r2": {
+        "unit": "kN/(km/h)^2",
+        "value": 0.00093264
+    },
+    "efficiency traction": {
+        "unit": "%",
+        "value": 0.9
+    },
+    "efficiency reg brake": {
+        "unit": "%",
+        "value": 0.9
+    },
+    "tunnel resistance": {
+        "units": {
+            "cross section": "m^2",
+            "resistance (or similar)": "kg/m"
+        },
+        "values": [
+            [
+                24,
+                15.16
+            ],
+            [
+                40,
+                7.23
+            ]
+        ]
+    }
+}
diff --git a/unitTests/trainLengthDependentTrackAttributes/speedLimit.py b/unitTests/trainLengthDependentTrackAttributes/speedLimit.py
new file mode 100644
index 0000000..e69de29
diff --git a/unitTests/tunnelResistance/tunnelResistance.py b/unitTests/tunnelResistance/tunnelResistance.py
new file mode 100644
index 0000000..abc1560
--- /dev/null
+++ b/unitTests/tunnelResistance/tunnelResistance.py
@@ -0,0 +1,18 @@
+import unittest
+
+from track import Track
+from train import Train
+
+
+class TestTunnelResistance(unittest.TestCase):
+
+    def testHigherEnergyConsumptionInTunnels(self):
+        '''
+
+        '''
+
+        train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='../trains')
+
+        track = Track(config={'id': 'CH_ZH_LU'}, pathJSON='../tracks')
+
+        self.assertTrue()
\ No newline at end of file

From 6a0b452f1d13b76229aca0271d0dc885d62a4df1 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Tue, 26 May 2026 15:51:57 +0200
Subject: [PATCH 14/31] test tunnel resistance added

---
 ocp.py                                        | 13 +++---
 track.py                                      | 42 +++++++++--------
 tracks/test_flat_no_tunnel.json               | 31 +++++++++++++
 tracks/test_flat_with_tunnel.json             | 45 +++++++++++++++++++
 .../tunnelResistance/tunnelResistance.py      | 31 +++++++++++--
 5 files changed, 135 insertions(+), 27 deletions(-)
 create mode 100644 tracks/test_flat_no_tunnel.json
 create mode 100644 tracks/test_flat_with_tunnel.json

diff --git a/ocp.py b/ocp.py
index 43b80b6..d66f0be 100644
--- a/ocp.py
+++ b/ocp.py
@@ -27,6 +27,8 @@ def __init__(self, paramsDict):
 
         self.integrateLosses = False  # integrate losses or take mid-point rule
 
+        self.withTrainLengthDependentTrackAttributes = False
+
         self.chooseClosestTunnelCrossSection = True # if exact tunnel cross section is not specified in train tunnel resistances, choose the closest value
 
         super().__init__(paramsDict)
@@ -123,7 +125,7 @@ def __init__(self, train, track, optsDict={}):
 
         # track parameters
 
-        self.points = computeDiscretizationPoints(track, numIntervals)
+        self.points = computeDiscretizationPoints(track, numIntervals, opts)
         self.steps = np.diff(self.points.index)
 
         # real-time parameters
@@ -195,12 +197,12 @@ def __init__(self, train, track, optsDict={}):
 
                 # gradient and curvature of current index
                 grad = self.points.iloc[i]['Gradient [permil]']/1e3
-                # gradLinearTerm = 0
-                gradLinearTerm = self.points.iloc[i]["Gradient linear term [permil]"]/1e3
+                gradLinearTerm = self.points.iloc[i]["Gradient linear term [permil/m]"]/1e3
                 curv = self.points.iloc[i]['Curvature [1/m]']
                 curvLinearTerm = self.points.iloc[i]["Curvature linear term [1/m^2]"]
                 crossSection = self.points.iloc[i]['CrossSection [m^2]']
                 tunnelFactor = computeTunnelFactor(crossSection, train, opts)
+                print(tunnelFactor)
 
                 # acceleration constraints
                 g += [trainModel.accelerationFun(ca.vertcat(time[i], velSq[i], 0), ca.vcat(u), grad, gradLinearTerm, curv, curvLinearTerm, tunnelFactor)]
@@ -209,7 +211,8 @@ def __init__(self, train, track, optsDict={}):
 
                 # coupling constraints
                 out = trainIntegrator.solve(time=time[i], velocitySquared=velSq[i], ds=self.steps[i],
-                    traction=Fel[i], pnBrake=Fpb[i], gradient=grad, curvature=curv)
+                    traction=Fel[i], pnBrake=Fpb[i], gradient=grad, gradientLinearTerm=curvLinearTerm, curvature=curv,
+                                            curvatureLinearTerm=curvLinearTerm, tunnelFactor=tunnelFactor)
 
                 xNxt1 = ca.vertcat(time[i+1], velSq[i+1])
                 xNxt2 = ca.vertcat(out['time'], out['velSquared'])
@@ -329,7 +332,7 @@ def solve(self, terminalTime, initialTime=0, terminalVelocity=1, initialVelocity
         # initial guess
         # NOTE: good idea vel0 to be compatible with f0 (power-wise) to avoid nans at first iteration
 
-        vel_avg = (self.points.index[-1]-self.points.index[0])/(terminalTime-initialTime)
+        vel_avg = (self.points.index[-1]-self.points.index[0])/(terminalTime-initialTime) * 0.95
         vel0 = vel_avg*vel_avg
         dt = (terminalTime - initialTime)/self.numIntervals
         t0 = initialTime
diff --git a/track.py b/track.py
index 3c299d9..d2b0590 100644
--- a/track.py
+++ b/track.py
@@ -89,7 +89,7 @@ def computeAltitude(gradients, length, altitudeStart=0):
     return df
 
 
-def computeDiscretizationPoints(track, numIntervals):
+def computeDiscretizationPoints(track, numIntervals, opts):
     """
     Compute the space discretization points based on track characteristics and horizon length.
     """
@@ -105,26 +105,32 @@ def computeDiscretizationPoints(track, numIntervals):
 
         raise ValueError("Wrong number of computed discretization intervals!")
 
-    # adapt constant track attribute terms
+    if opts.withTrainLengthDependentTrackAttributes:
+        # adapt constant track attribute terms to new shooting nodes
 
-    positions = df3.index.to_numpy(dtype=float)
-    grads = [df3["Gradient [permil]"].iloc[0]]
-    curvs = [df3["Curvature [1/m]"].iloc[0]]
+        positions = df3.index.to_numpy(dtype=float)
+        grads = [df3["Gradient [permil]"].iloc[0]]
+        curvs = [df3["Curvature [1/m]"].iloc[0]]
 
-    for idx in range(1, numIntervals+1):
+        for idx in range(1, numIntervals+1):
 
-        if np.isclose(df3["Gradient [permil]"].iloc[idx-1], df3["Gradient [permil]"].iloc[idx]):
-            grads.append(grads[-1] + (positions[idx] -  positions[idx - 1]) * df3["Gradient linear term [permil]"].iloc[idx - 1])
-        else:
-            grads.append(df3["Gradient [permil]"].iloc[idx])
+            if np.isclose(df3["Gradient [permil]"].iloc[idx-1], df3["Gradient [permil]"].iloc[idx]):
+                grads.append(grads[-1] + (positions[idx] -  positions[idx - 1]) * df3["Gradient linear term [permil/m]"].iloc[idx - 1])
+            else:
+                grads.append(df3["Gradient [permil]"].iloc[idx])
+
+            if np.isclose(df3["Curvature [1/m]"].iloc[idx - 1], df3["Curvature [1/m]"].iloc[idx]):
+                curvs.append(curvs[-1] + (positions[idx] - positions[idx - 1]) * df3["Curvature linear term [1/m^2]"].iloc[idx - 1])
+            else:
+                curvs.append(df3["Curvature [1/m]"].iloc[idx])
+
+        df3["Gradient [permil]"] = grads
+        df3["Curvature [1/m]"] = curvs
 
-        if np.isclose(df3["Curvature [1/m]"].iloc[idx - 1], df3["Curvature [1/m]"].iloc[idx]):
-            curvs.append(curvs[-1] + (positions[idx] - positions[idx - 1]) * df3["Curvature linear term [1/m^2]"].iloc[idx - 1])
-        else:
-            curvs.append(df3["Curvature [1/m]"].iloc[idx])
+    else:
 
-    df3["Gradient [permil]"] = grads
-    df3["Curvature [1/m]"] = curvs
+        df3["Gradient linear term [permil/m]"] =  np.zeros(len(df3))
+        df3["Curvature linear term [1/m^2]"] =  np.zeros(len(df3))
 
     return df3
 
@@ -154,7 +160,7 @@ def __init__(self, config, pathJSON='tracks'):
 
             data = json.load(file)
 
-        checkTTOBenchVersion(data, ['1.1', '1.2', '1.3'])
+        checkTTOBenchVersion(data, ['1.1', '1.2', '1.3', '1.4'])
 
         # read data
         self.length = convertUnit(data['stops']['values'][-1], data['stops']['unit'])
@@ -173,7 +179,7 @@ def __init__(self, config, pathJSON='tracks'):
 
         self.importTunnelTuples(data['tunnels']['values'] if 'tunnels' in data else [(0.0, 0.0, "infinity")],
                                 data['tunnels']['units']['length'] if 'tunnels' in data else 'm',
-                                data['tunnels']['units']['cross_section'] if 'tunnels' in data else 'm^2')
+                                data['tunnels']['units']['cross section'] if 'tunnels' in data else 'm^2')
 
 
         numStops = len(data['stops']['values'])
diff --git a/tracks/test_flat_no_tunnel.json b/tracks/test_flat_no_tunnel.json
new file mode 100644
index 0000000..1e47af0
--- /dev/null
+++ b/tracks/test_flat_no_tunnel.json
@@ -0,0 +1,31 @@
+{
+    "metadata": {
+        "id": "test_flat_no_tunnel",
+        "created by": "Roland Staerk",
+        "library version": "TTOBench v1.4",
+        "license": "BSD 2-Clause License"
+    },
+    "altitude": {
+        "unit": "m",
+        "value": 0
+    },
+    "stops": {
+        "unit": "m",
+        "values": [
+            0.0,
+            30000.0
+        ]
+    },
+    "speed limits": {
+        "units": {
+            "position": "m",
+            "velocity": "km/h"
+        },
+        "values": [
+            [
+                0.0,
+                160
+            ]
+        ]
+    }
+}
\ No newline at end of file
diff --git a/tracks/test_flat_with_tunnel.json b/tracks/test_flat_with_tunnel.json
new file mode 100644
index 0000000..1a40288
--- /dev/null
+++ b/tracks/test_flat_with_tunnel.json
@@ -0,0 +1,45 @@
+{
+    "metadata": {
+        "id": "test_flat_with_tunnel",
+        "created by": "Roland Staerk",
+        "library version": "TTOBench v1.4",
+        "license": "BSD 2-Clause License"
+    },
+    "altitude": {
+        "unit": "m",
+        "value": 0
+    },
+    "stops": {
+        "unit": "m",
+        "values": [
+            0.0,
+            30000.0
+        ]
+    },
+    "speed limits": {
+        "units": {
+            "position": "m",
+            "velocity": "km/h"
+        },
+        "values": [
+            [
+                0.0,
+                160
+            ]
+        ]
+    },
+    "tunnels": {
+        "units": {
+            "position": "m",
+            "length": "m",
+            "cross section": "m^2"
+        },
+        "values": [
+            [
+                1000.0,
+                26000.0,
+                24.0
+            ]
+        ]
+    }
+}
\ No newline at end of file
diff --git a/unitTests/tunnelResistance/tunnelResistance.py b/unitTests/tunnelResistance/tunnelResistance.py
index abc1560..564f986 100644
--- a/unitTests/tunnelResistance/tunnelResistance.py
+++ b/unitTests/tunnelResistance/tunnelResistance.py
@@ -1,5 +1,8 @@
+import os
 import unittest
+from pathlib import Path
 
+from ocp import casadiSolver
 from track import Track
 from train import Train
 
@@ -8,11 +11,31 @@ class TestTunnelResistance(unittest.TestCase):
 
     def testHigherEnergyConsumptionInTunnels(self):
         '''
-
+        26 km long small tunnel with cross section of 24 m^2 on a track of 28 km results in significant higher energy consumption.
         '''
 
-        train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='../trains')
+        startPosition = 0  # [m]
+        endPosition = 28000  # [m]
+        duration = 28000/(145/3.6)  # [s]
+
+        train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='trains')
+
+        track = Track(config={'id': 'test_flat_no_tunnel'}, pathJSON='tracks')
+        track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
+
+        opts = {'numIntervals': 300, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True}
+        solver = casadiSolver(train, track, opts)
+        df1, stats1 = solver.solve(duration)
+
+        energyConsuptionWithoutTunnel = stats1['Cost']
+
+        track = Track(config={'id': 'test_flat_with_tunnel'}, pathJSON='tracks')
+        track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
+
+        opts = {'numIntervals': 300, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True}
+        solver = casadiSolver(train, track, opts)
+        df2, stats2 = solver.solve(duration)
 
-        track = Track(config={'id': 'CH_ZH_LU'}, pathJSON='../tracks')
+        energyConsuptionWithTunnel = stats2['Cost']
 
-        self.assertTrue()
\ No newline at end of file
+        self.assertTrue(energyConsuptionWithoutTunnel < energyConsuptionWithTunnel)
\ No newline at end of file

From 97c20ab4be932235eda25845c0e683f668a52059 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Wed, 27 May 2026 09:32:26 +0200
Subject: [PATCH 15/31] test speed limit added

---
 ocp.py                                        |  3 -
 track.py                                      |  2 +-
 tracks/test_speed_increase.json               | 35 +++++++++++
 .../speedLimit.py                             | 61 +++++++++++++++++++
 .../tunnelResistance/tunnelResistance.py      | 31 +++++++---
 5 files changed, 120 insertions(+), 12 deletions(-)
 create mode 100644 tracks/test_speed_increase.json

diff --git a/ocp.py b/ocp.py
index d66f0be..7d7e09c 100644
--- a/ocp.py
+++ b/ocp.py
@@ -27,8 +27,6 @@ def __init__(self, paramsDict):
 
         self.integrateLosses = False  # integrate losses or take mid-point rule
 
-        self.withTrainLengthDependentTrackAttributes = False
-
         self.chooseClosestTunnelCrossSection = True # if exact tunnel cross section is not specified in train tunnel resistances, choose the closest value
 
         super().__init__(paramsDict)
@@ -202,7 +200,6 @@ def __init__(self, train, track, optsDict={}):
                 curvLinearTerm = self.points.iloc[i]["Curvature linear term [1/m^2]"]
                 crossSection = self.points.iloc[i]['CrossSection [m^2]']
                 tunnelFactor = computeTunnelFactor(crossSection, train, opts)
-                print(tunnelFactor)
 
                 # acceleration constraints
                 g += [trainModel.accelerationFun(ca.vertcat(time[i], velSq[i], 0), ca.vcat(u), grad, gradLinearTerm, curv, curvLinearTerm, tunnelFactor)]
diff --git a/track.py b/track.py
index d2b0590..f4cc0bf 100644
--- a/track.py
+++ b/track.py
@@ -105,7 +105,7 @@ def computeDiscretizationPoints(track, numIntervals, opts):
 
         raise ValueError("Wrong number of computed discretization intervals!")
 
-    if opts.withTrainLengthDependentTrackAttributes:
+    if "Gradient linear term [permil/m]" in df3.columns:
         # adapt constant track attribute terms to new shooting nodes
 
         positions = df3.index.to_numpy(dtype=float)
diff --git a/tracks/test_speed_increase.json b/tracks/test_speed_increase.json
new file mode 100644
index 0000000..8b53643
--- /dev/null
+++ b/tracks/test_speed_increase.json
@@ -0,0 +1,35 @@
+{
+    "metadata": {
+        "id": "test_flat_with_tunnel",
+        "created by": "Roland Staerk",
+        "library version": "TTOBench v1.4",
+        "license": "BSD 2-Clause License"
+    },
+    "altitude": {
+        "unit": "m",
+        "value": 0
+    },
+    "stops": {
+        "unit": "m",
+        "values": [
+            0.0,
+            12000.0
+        ]
+    },
+    "speed limits": {
+        "units": {
+            "position": "m",
+            "velocity": "km/h"
+        },
+        "values": [
+            [
+                0.0,
+                79.2
+            ],
+            [
+                1000.0,
+                144
+            ]
+        ]
+    }
+}
\ No newline at end of file
diff --git a/unitTests/trainLengthDependentTrackAttributes/speedLimit.py b/unitTests/trainLengthDependentTrackAttributes/speedLimit.py
index e69de29..68fba47 100644
--- a/unitTests/trainLengthDependentTrackAttributes/speedLimit.py
+++ b/unitTests/trainLengthDependentTrackAttributes/speedLimit.py
@@ -0,0 +1,61 @@
+import unittest
+
+from ocp import casadiSolver
+from track import Track
+from train import Train
+
+
+class TestSpeedLimit(unittest.TestCase):
+
+    def testTrainDoesNotAccelerateToEarly(self):
+        '''
+        Speed limit increases from 22 m/s to 40 m/s at position 1000 m.
+
+        Without train-length-dependent values, the optimizer may accelerate too early.
+        With train-length-dependent values, the front of the train may only exceed
+        22 m/s after the whole train has passed the speed-increase position.
+        '''
+
+        startPosition = 0 # [m]
+        endPosition = 12000 # [m]
+        duration = 12000/(115/3.6) # [s]
+
+        train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='trains')
+
+        track = Track(config={'id': 'test_speed_increase'}, pathJSON='tracks')
+        track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
+
+        opts = {'numIntervals': 300, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True}
+        solver = casadiSolver(train, track, opts)
+        df1, stats1 = solver.solve(duration)
+
+        speedAfterSpeedIncrease1 = df1[df1['Position [m]'] > 1000].iloc[0]['Velocity [m/s]']
+
+        track.updateTrainLengthDependentValues(train)
+
+        opts = {'numIntervals': 300, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True}
+        solver = casadiSolver(train, track, opts)
+        df2, stats2 = solver.solve(duration)
+
+        speedAfterSpeedIncrease2 = df2[df2['Position [m]'] > 1000].iloc[0]['Velocity [m/s]']
+        speedAfterSpeedIncrease3 = df2[df2['Position [m]'] > (1000+train.length)].iloc[0]['Velocity [m/s]']
+
+        epsilon = 0.001
+
+        self.assertGreater(
+            speedAfterSpeedIncrease1,
+            22,
+            msg="Without train-length-dependent values, the train should accelerate immediately after the speed limit increase."
+        )
+
+        self.assertLessEqual(
+            speedAfterSpeedIncrease2,
+            22 + epsilon,
+            msg="With train-length-dependent values, the train should not accelerate before the whole train has passed the speed limit increase."
+        )
+
+        self.assertGreater(
+            speedAfterSpeedIncrease3,
+            22,
+            msg="With train-length-dependent values, the train should accelerate after the whole train has passed the speed limit increase."
+        )
\ No newline at end of file
diff --git a/unitTests/tunnelResistance/tunnelResistance.py b/unitTests/tunnelResistance/tunnelResistance.py
index 564f986..83df459 100644
--- a/unitTests/tunnelResistance/tunnelResistance.py
+++ b/unitTests/tunnelResistance/tunnelResistance.py
@@ -1,6 +1,4 @@
-import os
 import unittest
-from pathlib import Path
 
 from ocp import casadiSolver
 from track import Track
@@ -14,9 +12,9 @@ def testHigherEnergyConsumptionInTunnels(self):
         26 km long small tunnel with cross section of 24 m^2 on a track of 28 km results in significant higher energy consumption.
         '''
 
-        startPosition = 0  # [m]
-        endPosition = 28000  # [m]
-        duration = 28000/(145/3.6)  # [s]
+        startPosition = 0 # [m]
+        endPosition = 28000 # [m]
+        duration = 28000/(145/3.6) # [s]
 
         train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='trains')
 
@@ -27,7 +25,7 @@ def testHigherEnergyConsumptionInTunnels(self):
         solver = casadiSolver(train, track, opts)
         df1, stats1 = solver.solve(duration)
 
-        energyConsuptionWithoutTunnel = stats1['Cost']
+        energyConsumptionWithoutTunnel = stats1['Cost']
 
         track = Track(config={'id': 'test_flat_with_tunnel'}, pathJSON='tracks')
         track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
@@ -36,6 +34,23 @@ def testHigherEnergyConsumptionInTunnels(self):
         solver = casadiSolver(train, track, opts)
         df2, stats2 = solver.solve(duration)
 
-        energyConsuptionWithTunnel = stats2['Cost']
+        energyConsumptionWithTunnel = stats2['Cost']
 
-        self.assertTrue(energyConsuptionWithoutTunnel < energyConsuptionWithTunnel)
\ No newline at end of file
+        minEnergyRatio = 1.5
+
+        self.assertGreater(
+            energyConsumptionWithTunnel,
+            energyConsumptionWithoutTunnel,
+            msg="Energy consumption with tunnel should be higher than without tunnel."
+        )
+
+        self.assertGreater(
+            energyConsumptionWithTunnel / energyConsumptionWithoutTunnel,
+            minEnergyRatio,
+            msg=(
+                "Energy consumption with tunnel should be significantly higher. "
+                f"Expected ratio > {minEnergyRatio}, got "
+                f"{energyConsumptionWithTunnel / energyConsumptionWithoutTunnel:.3f}."
+            )
+
+        )
\ No newline at end of file

From 1a2a00b0bdc29a5712225431bf62b972344a4569 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Wed, 27 May 2026 14:32:10 +0200
Subject: [PATCH 16/31] test gradients added

---
 ocp.py                                        |  2 +
 track.py                                      | 71 ++++++++++---
 tracks/test_one_hill.json                     | 59 +++++++++++
 ...ease.json => test_one_speed_increase.json} |  2 +-
 train.py                                      | 12 +--
 .../gradient.py                               | 99 +++++++++++++++++++
 utils.py                                      |  7 ++
 7 files changed, 233 insertions(+), 19 deletions(-)
 create mode 100644 tracks/test_one_hill.json
 rename tracks/{test_speed_increase.json => test_one_speed_increase.json} (93%)
 create mode 100644 unitTests/trainLengthDependentTrackAttributes/gradient.py

diff --git a/ocp.py b/ocp.py
index 7d7e09c..b76cdd6 100644
--- a/ocp.py
+++ b/ocp.py
@@ -29,6 +29,8 @@ def __init__(self, paramsDict):
 
         self.chooseClosestTunnelCrossSection = True # if exact tunnel cross section is not specified in train tunnel resistances, choose the closest value
 
+        self.pwcLengthDependentTrackAttributes = False
+
         super().__init__(paramsDict)
 
 
diff --git a/track.py b/track.py
index f4cc0bf..0d0cb30 100644
--- a/track.py
+++ b/track.py
@@ -105,34 +105,81 @@ def computeDiscretizationPoints(track, numIntervals, opts):
 
         raise ValueError("Wrong number of computed discretization intervals!")
 
+    adaptConstantTrackAttributesToNewShootingNodes(df3, numIntervals)
+
+    if opts.pwcLengthDependentTrackAttributes:
+
+        makePwcLengthDependentTrackAttibutes(df3)
+
+    return df3
+
+
+def adaptConstantTrackAttributesToNewShootingNodes(df3, numIntervals):
+
     if "Gradient linear term [permil/m]" in df3.columns:
-        # adapt constant track attribute terms to new shooting nodes
 
         positions = df3.index.to_numpy(dtype=float)
         grads = [df3["Gradient [permil]"].iloc[0]]
-        curvs = [df3["Curvature [1/m]"].iloc[0]]
 
-        for idx in range(1, numIntervals+1):
+        for idx in range(1, numIntervals + 1):
+
+            if np.isclose(df3["Gradient [permil]"].iloc[idx - 1], df3["Gradient [permil]"].iloc[idx]):
+
+                grads.append(grads[-1] + (positions[idx] - positions[idx - 1]) * df3["Gradient linear term [permil/m]"].iloc[idx - 1])
 
-            if np.isclose(df3["Gradient [permil]"].iloc[idx-1], df3["Gradient [permil]"].iloc[idx]):
-                grads.append(grads[-1] + (positions[idx] -  positions[idx - 1]) * df3["Gradient linear term [permil/m]"].iloc[idx - 1])
             else:
+
                 grads.append(df3["Gradient [permil]"].iloc[idx])
 
+        df3["Gradient [permil]"] = grads
+
+    else:
+
+        df3["Gradient linear term [permil/m]"] = np.zeros(len(df3))
+
+
+    if "Curvature linear term [1/m^2]" in df3.columns:
+
+        positions = df3.index.to_numpy(dtype=float)
+        curvs = [df3["Curvature [1/m]"].iloc[0]]
+
+        for idx in range(1, numIntervals + 1):
+
             if np.isclose(df3["Curvature [1/m]"].iloc[idx - 1], df3["Curvature [1/m]"].iloc[idx]):
+
                 curvs.append(curvs[-1] + (positions[idx] - positions[idx - 1]) * df3["Curvature linear term [1/m^2]"].iloc[idx - 1])
+
             else:
+
                 curvs.append(df3["Curvature [1/m]"].iloc[idx])
 
-        df3["Gradient [permil]"] = grads
         df3["Curvature [1/m]"] = curvs
 
     else:
 
-        df3["Gradient linear term [permil/m]"] =  np.zeros(len(df3))
-        df3["Curvature linear term [1/m^2]"] =  np.zeros(len(df3))
+        df3["Curvature linear term [1/m^2]"] = np.zeros(len(df3))
 
-    return df3
+
+def makePwcLengthDependentTrackAttibutes(df3):
+
+    positions = df3.index.to_numpy(dtype=float)
+
+    g_pwl = df3["Gradient [permil]"].to_numpy(dtype=float)
+    g_linear = df3["Gradient linear term [permil/m]"].to_numpy(dtype=float)
+
+    c_pwl = df3["Curvature [1/m]"].to_numpy(dtype=float)
+    c_linear = df3["Curvature linear term [1/m^2]"].to_numpy(dtype=float)
+
+    ds = positions[1:] - positions[:-1]
+
+    g_pwc = g_pwl[:-1] + 0.5 * g_linear[:-1] * ds
+    c_pwc = c_pwl[:-1] + 0.5 * c_linear[:-1] * ds
+
+    df3["Gradient [permil]"] = np.r_[g_pwc, g_pwc[-1]]
+    df3["Curvature [1/m]"] = np.r_[c_pwc, c_pwc[-1]]
+
+    df3["Gradient linear term [permil/m]"] = np.zeros(len(df3))
+    df3["Curvature linear term [1/m^2]"] = np.zeros(len(df3))
 
 
 class Track():
@@ -590,7 +637,7 @@ def updateGradientsToTrainLength(self, trainLength):
 
         if len(pos) > 1:
 
-            pos_adj = np.sort(np.r_[pos, pos + trainLength])
+            pos_adj = np.sort(np.unique(np.r_[pos, pos + trainLength]))
 
             pos_adj = pos_adj[pos_adj < self.length]
 
@@ -607,7 +654,7 @@ def updateGradientsToTrainLength(self, trainLength):
                 epsilon = 0.001
                 list_indices = []
 
-                for idx2 in range(len(pos)-1):
+                for idx2 in range(len(pos)):
 
                     if currentPosition - trainLength + epsilon < pos[idx2] < currentPosition + epsilon:
                         list_indices.append(idx2)
@@ -620,7 +667,7 @@ def updateGradientsToTrainLength(self, trainLength):
             # plotGradients(self, np.asarray(pos_adj, dtype=float), np.asarray(g_adj, dtype=float), np.asarray(g_linear, dtype=float))
 
             self.gradients = pd.DataFrame(
-                {"Gradient [permil]": g_adj, "Gradient linear term [permil]": g_linear},
+                {"Gradient [permil]": g_adj, "Gradient linear term [permil/m]": g_linear},
                 index=pos_adj,
             )
 
diff --git a/tracks/test_one_hill.json b/tracks/test_one_hill.json
new file mode 100644
index 0000000..42f751a
--- /dev/null
+++ b/tracks/test_one_hill.json
@@ -0,0 +1,59 @@
+{
+    "metadata": {
+        "id": "test_one_hill",
+        "created by": "Roland Staerk",
+        "library version": "TTOBench v1.4",
+        "license": "BSD 2-Clause License"
+    },
+    "altitude": {
+        "unit": "m",
+        "value": 0
+    },
+    "stops": {
+        "unit": "m",
+        "values": [
+            0.0,
+            12000.0
+        ]
+    },
+    "speed limits": {
+        "units": {
+            "position": "m",
+            "velocity": "km/h"
+        },
+        "values": [
+            [
+                0.0,
+                79.2
+            ]
+        ]
+    },
+    "gradients": {
+        "units": {
+            "position": "m",
+            "slope": "permil"
+        },
+        "values": [
+            [
+                0.0,
+                0.0
+            ],
+            [
+                1000.0,
+                20.0
+            ],
+            [
+                2000.0,
+                0.0
+            ],
+            [
+                3000.0,
+                -20.0
+            ],
+            [
+                4000.0,
+                0.0
+            ]
+        ]
+    }
+}
\ No newline at end of file
diff --git a/tracks/test_speed_increase.json b/tracks/test_one_speed_increase.json
similarity index 93%
rename from tracks/test_speed_increase.json
rename to tracks/test_one_speed_increase.json
index 8b53643..dbc4a5e 100644
--- a/tracks/test_speed_increase.json
+++ b/tracks/test_one_speed_increase.json
@@ -1,6 +1,6 @@
 {
     "metadata": {
-        "id": "test_flat_with_tunnel",
+        "id": "test_one_speed_increase",
         "created by": "Roland Staerk",
         "library version": "TTOBench v1.4",
         "license": "BSD 2-Clause License"
diff --git a/train.py b/train.py
index 37b92af..fbfb2c7 100644
--- a/train.py
+++ b/train.py
@@ -264,7 +264,7 @@ def __init__(self, sr0, sr1, sr2, rho=1, g=9.81, withPnBrake=True) -> None:
         gradient = ca.MX.sym('gradient')                        # [-]     -> values between [0,0.2]
         gradientLinearTerm = ca.MX.sym('gradientLinearTerm')    # [1/m]
         curvature = ca.MX.sym('curvature')                      # [1/m]   -> values between [0,0.004]
-        curvatureLinearTerm = ca.MX.sym('curvatureLinearTerm')            # [1/m^2]
+        curvatureLinearTerm = ca.MX.sym('curvatureLinearTerm')  # [1/m^2]
         tunnelFactor = ca.MX.sym('tunnelFactor')                # [1/m]
         ds = ca.MX.sym('ds')
 
@@ -272,10 +272,10 @@ def __init__(self, sr0, sr1, sr2, rho=1, g=9.81, withPnBrake=True) -> None:
 
         # ODE
 
-        rollingResistance = sr0 + sr1*ca.sqrt(velocitySquared) + sr2*velocitySquared                # [N/kg]
-        gradientResistance = g*(1/rho)*(gradient+gradientLinearTerm*position)                                                  # [N/kg]
-        curvatureResistance = (1/rho)*(5.07 * (curvature+curvatureLinearTerm*position))                                                          # [N/kg]
-        tunnelResistance = tunnelFactor * velocitySquared                                           # [N/kg]
+        rollingResistance = sr0 + sr1*ca.sqrt(velocitySquared) + sr2*velocitySquared # [N/kg]
+        gradientResistance = g*(1/rho)*(gradient+gradientLinearTerm*position) # [N/kg]
+        curvatureResistance = (1/rho)*(5.07 * (curvature+curvatureLinearTerm*position)) # [N/kg]
+        tunnelResistance = tunnelFactor * velocitySquared # [N/kg]
 
         acceleration = traction + (pnBrake if withPnBrake else 0) - rollingResistance - gradientResistance - curvatureResistance - tunnelResistance # [m/s^2]
 
@@ -378,7 +378,7 @@ def __init__(self, model, solver, optsDict={}) -> None:
 
             eval = ca.vertcat(tApprox, zf[0, ns], zf[1, ns])
 
-            self.eval = ca.Function('xNxt', [model.states, ca.vertcat(model.controls, model.parameters), ca.MX.sym('ds')], [eval])
+            self.eval = ca.Function('xNxt', [model.states, ca.vertcat(model.controls, model.parameters), ca.MX.sym('dummy')], [eval])
 
 
     def solve(self, time, velocitySquared, ds, position=0, traction=0, pnBrake=0, gradient=0, gradientLinearTerm=0, curvature=0, curvatureLinearTerm=0, tunnelFactor=0):
diff --git a/unitTests/trainLengthDependentTrackAttributes/gradient.py b/unitTests/trainLengthDependentTrackAttributes/gradient.py
new file mode 100644
index 0000000..97cc11e
--- /dev/null
+++ b/unitTests/trainLengthDependentTrackAttributes/gradient.py
@@ -0,0 +1,99 @@
+import unittest
+
+from matplotlib import pyplot as plt
+
+from ocp import casadiSolver
+from track import Track, computeAltitude
+from train import Train
+
+
+class TestGradient(unittest.TestCase):
+
+    def testLinearGradient(self):
+        '''
+        Track with 20 permil increase from 1000 m to 2000 m and 20 permil
+        decrease from 3000 m to 4000 m.
+
+        Energy consumption should be roughly equal if computed using piecewise
+        linear gradients or equivalent piecewise constant gradients.
+
+        Altitude should be 0 m at the end.
+        '''
+
+        startPosition = 0 # [m]
+        endPosition = 5000 # [m]
+        duration = 5000/(60/3.6) # [s]
+
+        altitudeTolerance = 1e-4
+        energyRelativeTolerance = 0.004
+        numIntervals = 100
+
+        train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='trains')
+        train.length = 600
+
+        track = Track(config={'id': 'test_one_hill'}, pathJSON='tracks')
+        track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
+        track.updateTrainLengthDependentValues(train)
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True}
+        solver = casadiSolver(train, track, opts)
+        df1, stats1 = solver.solve(duration)
+
+        energyConsumptionWithLinearTerms = stats1['Cost']
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True, 'pwcLengthDependentTrackAttributes': True}
+        solver = casadiSolver(train, track, opts)
+        df2, stats2 = solver.solve(duration)
+
+        energyConsumptionWithPwcTerms= stats2['Cost']
+
+        relativeEnergyDifference = (abs(energyConsumptionWithLinearTerms - energyConsumptionWithPwcTerms) / energyConsumptionWithLinearTerms)
+
+        df_grads_2 = df2.set_index("Position [m]")[["Gradient [permil]"]]
+        df_alt_2 = computeAltitude(df_grads_2, track.length)
+        finalAltitude = df_alt_2.iloc[-1]["Altitude [m]"]
+
+        self.assertLess(
+            relativeEnergyDifference,
+            energyRelativeTolerance,
+            msg=(
+                "Energy consumption differs too much between piecewise linear and "
+                f"piecewise constant gradients. Relative difference: "
+                f"{relativeEnergyDifference:.6f}, tolerance: {energyRelativeTolerance:.6f}. "
+                f"PWL cost: {energyConsumptionWithLinearTerms:.6f}, "
+                f"PWC cost: {energyConsumptionWithPwcTerms:.6f}."
+            )
+        )
+
+        self.assertLess(
+            abs(finalAltitude),
+            altitudeTolerance,
+            msg=(
+                "Final altitude should be close to 0 m. "
+                f"Final altitude: {finalAltitude:.8f} m, "
+                f"tolerance: {altitudeTolerance:.8f} m."
+            )
+        )
+
+        plotDebug = True
+
+        if plotDebug:
+
+            fig, ax = plt.subplots(figsize=(16, 8))
+
+            df_grads_1 = df1.set_index("Position [m]")[["Gradient [permil]"]]
+            ax.plot(df_grads_1.index.values / 1000, df_grads_1["Gradient [permil]"],label="pwl gradients")
+            ax.step(df_grads_2.index.values / 1000, df_grads_2["Gradient [permil]"], '--', where='post', label="pwc gradients")
+            ax.set_title("Gradients")
+            ax.set_xlabel("Position [km]")
+            ax.set_ylabel("Gradient [‰]")
+            ax.grid(True, which="both", linestyle="--", alpha=0.5)
+            ax.legend(loc="upper right")
+            ax.set_xlim(0, track.length / 1000)
+            ax.figure.tight_layout()
+
+            plt.show()
+
+    def testAllIntegratorsWork(self):
+        startPosition = 0
+        endPosition = 12000
\ No newline at end of file
diff --git a/utils.py b/utils.py
index c93cf8c..fd00f80 100644
--- a/utils.py
+++ b/utils.py
@@ -578,6 +578,13 @@ def plotGradients(track, pos_adj, g_adj, g_linear):
     x = track.gradients.index.to_numpy(dtype=float)
     g = track.gradients["Gradient [permil]"].to_numpy(dtype=float)
 
+    x = np.append(x, track.length)
+    g = np.append(g, g[-1])
+
+    pos_adj = np.append(pos_adj, track.length)
+    g_adj = np.append(g_adj, g_adj[-1])
+    g_linear = np.append(g_linear, g_linear[-1])
+
     fig, ax = plt.subplots(figsize=(16, 8))
 
     ax.step(x/1000, g, where='post', label="piecewise constant gradients")

From 8fe528e66365e861db752e94406fcaf9dfd4a911 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Wed, 27 May 2026 15:56:16 +0200
Subject: [PATCH 17/31] test curvature added

---
 track.py                                      |   2 +-
 tracks/test_two_radii.json                    |  65 +++++++
 .../curvature.py                              | 168 ++++++++++++++++++
 .../gradient.py                               |  93 +++++++++-
 4 files changed, 324 insertions(+), 4 deletions(-)
 create mode 100644 tracks/test_two_radii.json
 create mode 100644 unitTests/trainLengthDependentTrackAttributes/curvature.py

diff --git a/track.py b/track.py
index 0d0cb30..23392c8 100644
--- a/track.py
+++ b/track.py
@@ -696,7 +696,7 @@ def updateCurvaturesToTrainLength(self, trainLength):
                 epsilon = 0.001
                 list_indices = []
 
-                for idx2 in range(len(pos)-1):
+                for idx2 in range(len(pos)):
 
                     if currentPosition - trainLength + epsilon < pos[idx2] < currentPosition + epsilon:
                         list_indices.append(idx2)
diff --git a/tracks/test_two_radii.json b/tracks/test_two_radii.json
new file mode 100644
index 0000000..12b67c0
--- /dev/null
+++ b/tracks/test_two_radii.json
@@ -0,0 +1,65 @@
+{
+    "metadata": {
+        "id": "test_two_radii",
+        "created by": "Roland Staerk",
+        "library version": "TTOBench v1.4",
+        "license": "BSD 2-Clause License"
+    },
+    "altitude": {
+        "unit": "m",
+        "value": 0
+    },
+    "stops": {
+        "unit": "m",
+        "values": [
+            0.0,
+            12000.0
+        ]
+    },
+    "speed limits": {
+        "units": {
+            "position": "m",
+            "velocity": "km/h"
+        },
+        "values": [
+            [
+                0.0,
+                79.2
+            ]
+        ]
+    },
+    "curvatures": {
+        "units": {
+            "position": "m",
+            "radius at start": "m",
+            "radius at end": "m"
+        },
+        "values": [
+            [
+                0.0,
+                Infinity,
+                Infinity
+            ],
+            [
+                1000.0,
+                300.0,
+                300.0
+            ],
+            [
+                2000.0,
+                Infinity,
+                Infinity
+            ],
+            [
+                3000.0,
+                -300.0,
+                -300.0
+            ],
+            [
+                4000.0,
+                Infinity,
+                Infinity
+            ]
+        ]
+    }
+}
\ No newline at end of file
diff --git a/unitTests/trainLengthDependentTrackAttributes/curvature.py b/unitTests/trainLengthDependentTrackAttributes/curvature.py
new file mode 100644
index 0000000..4feb44e
--- /dev/null
+++ b/unitTests/trainLengthDependentTrackAttributes/curvature.py
@@ -0,0 +1,168 @@
+import unittest
+
+from matplotlib import pyplot as plt
+
+from ocp import casadiSolver
+from track import Track, computeAltitude
+from train import Train
+
+
+class TestCurvature(unittest.TestCase):
+
+    def testLinearCurvature(self):
+        '''
+        Track with right turn from 1000 m to 2000 m and a left turn from 3000 m to 4000 m.
+
+        Energy consumption should be roughly equal if computed using piecewise
+        linear curvatures or equivalent piecewise constant curvatures.
+        '''
+
+        startPosition = 0 # [m]
+        endPosition = 5000 # [m]
+        duration = 5000/(60/3.6) # [s]
+
+        energyRelativeTolerance = 0.004
+        numIntervals = 100
+
+        train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='trains')
+        train.length = 600
+
+        track = Track(config={'id': 'test_two_radii'}, pathJSON='tracks')
+        track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
+        track.updateTrainLengthDependentValues(train)
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True}
+        solver = casadiSolver(train, track, opts)
+        df1, stats1 = solver.solve(duration)
+
+        energyConsumptionWithLinearTerms = stats1['Cost']
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True, 'pwcLengthDependentTrackAttributes': True}
+        solver = casadiSolver(train, track, opts)
+        df2, stats2 = solver.solve(duration)
+
+        energyConsumptionWithPwcTerms= stats2['Cost']
+
+        relativeEnergyDifference = (abs(energyConsumptionWithLinearTerms - energyConsumptionWithPwcTerms) / energyConsumptionWithLinearTerms)
+
+        self.assertLess(
+            relativeEnergyDifference,
+            energyRelativeTolerance,
+            msg=(
+                "Energy consumption differs too much between piecewise linear and "
+                f"piecewise constant curvatures. Relative difference: "
+                f"{relativeEnergyDifference:.6f}, tolerance: {energyRelativeTolerance:.6f}. "
+                f"PWL cost: {energyConsumptionWithLinearTerms:.6f}, "
+                f"PWC cost: {energyConsumptionWithPwcTerms:.6f}."
+            )
+        )
+
+
+        plotDebug = True
+
+        if plotDebug:
+
+            fig, ax = plt.subplots(figsize=(16, 8))
+
+            ax.plot(df1["Position [m]"] / 1000, df1["Curvature [1/m]"], label="pwl curvatures")
+            ax.step(df2["Position [m]"] / 1000, df2["Curvature [1/m]"], "--", where="post", label="pwc curvatures")
+            ax.set_title("Curvatures")
+            ax.set_xlabel("Position [km]")
+            ax.set_ylabel("Curvature [1/m]")
+            ax.grid(True, which="both", linestyle="--", alpha=0.5)
+            ax.legend(loc="upper right")
+            ax.set_xlim(0, track.length / 1000)
+            ax.figure.tight_layout()
+
+            plt.show()
+
+    def testAllIntegratorTypesWork(self):
+        '''
+        Verify that all supported integration methods produce consistent results
+        for the same train, track, and optimization setup.
+
+        The test compares RK, IRK, and CVODES, including the approximate time
+        integration option for RK and IRK. The resulting energy costs should only
+        differ by a small relative tolerance.
+        '''
+
+        startPosition = 0  # [m]
+        endPosition = 5000  # [m]
+        duration = 5000 / (60 / 3.6)  # [s]
+
+        tol = 0.02
+        numIntervals = 100
+
+        train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='trains')
+        train.length = 600
+
+        track = Track(config={'id': 'test_two_radii'}, pathJSON='tracks')
+        track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
+        track.updateTrainLengthDependentValues(train)
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1},'energyOptimal': True}
+        solver = casadiSolver(train, track, opts)
+        df, stats = solver.solve(duration)
+
+        energy_RK_Approx = stats['Cost']
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 0},'energyOptimal': True}
+        solver = casadiSolver(train, track, opts)
+        df, stats = solver.solve(duration)
+
+        energy_RK = stats['Cost']
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'IRK', 'integrationOptions': {'numApproxSteps': 1},'energyOptimal': True}
+        solver = casadiSolver(train, track, opts)
+        df, stats = solver.solve(duration)
+
+        energy_IRK_Approx = stats['Cost']
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'IRK', 'integrationOptions': {'numApproxSteps': 0},'energyOptimal': True}
+        solver = casadiSolver(train, track, opts)
+        df, stats = solver.solve(duration)
+
+        energy_IRK = stats['Cost']
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'CVODES', 'energyOptimal': True}
+        solver = casadiSolver(train, track, opts)
+        df, stats = solver.solve(duration)
+
+        energy_CVODES = stats['Cost']
+
+        relDiff_RKApprox_IRKApprox = abs(energy_RK_Approx - energy_IRK_Approx) / energy_RK_Approx
+        relDiff_RKApprox_CVODES = abs(energy_RK_Approx - energy_CVODES) / energy_RK_Approx
+        relDiff_RK_IRK = abs(energy_RK - energy_IRK) / energy_RK
+
+        self.assertLess(
+            relDiff_RKApprox_IRKApprox,
+            tol,
+            msg=(
+                "RK and IRK with numApproxSteps=1 should give similar costs. "
+                f"RK approx: {energy_RK_Approx:.6f}, "
+                f"IRK approx: {energy_IRK_Approx:.6f}, "
+                f"relative difference: {relDiff_RKApprox_IRKApprox:.6f}."
+            )
+        )
+
+        self.assertLess(
+            relDiff_RKApprox_CVODES,
+            tol,
+            msg=(
+                "RK with numApproxSteps=1 and CVODES should give similar costs. "
+                f"RK approx: {energy_RK_Approx:.6f}, "
+                f"CVODES: {energy_CVODES:.6f}, "
+                f"relative difference: {relDiff_RKApprox_CVODES:.6f}."
+            )
+        )
+
+        self.assertLess(
+            relDiff_RK_IRK,
+            tol,
+            msg=(
+                "RK and IRK with numApproxSteps=0 should give similar costs. "
+                f"RK: {energy_RK:.6f}, "
+                f"IRK: {energy_IRK:.6f}, "
+                f"relative difference: {relDiff_RK_IRK:.6f}."
+            )
+        )
\ No newline at end of file
diff --git a/unitTests/trainLengthDependentTrackAttributes/gradient.py b/unitTests/trainLengthDependentTrackAttributes/gradient.py
index 97cc11e..7b089aa 100644
--- a/unitTests/trainLengthDependentTrackAttributes/gradient.py
+++ b/unitTests/trainLengthDependentTrackAttributes/gradient.py
@@ -94,6 +94,93 @@ def testLinearGradient(self):
 
             plt.show()
 
-    def testAllIntegratorsWork(self):
-        startPosition = 0
-        endPosition = 12000
\ No newline at end of file
+    def testAllIntegratorTypesWork(self):
+        '''
+        Verify that all supported integration methods produce consistent results
+        for the same train, track, and optimization setup.
+
+        The test compares RK, IRK, and CVODES, including the approximate time
+        integration option for RK and IRK. The resulting energy costs should only
+        differ by a small relative tolerance.
+        '''
+
+        startPosition = 0  # [m]
+        endPosition = 5000  # [m]
+        duration = 5000 / (60 / 3.6)  # [s]
+
+        tol = 0.02
+        numIntervals = 100
+
+        train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='trains')
+        train.length = 600
+
+        track = Track(config={'id': 'test_one_hill'}, pathJSON='tracks')
+        track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
+        track.updateTrainLengthDependentValues(train)
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1},'energyOptimal': True}
+        solver = casadiSolver(train, track, opts)
+        df, stats = solver.solve(duration)
+
+        energy_RK_Approx = stats['Cost']
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 0},'energyOptimal': True}
+        solver = casadiSolver(train, track, opts)
+        df, stats = solver.solve(duration)
+
+        energy_RK = stats['Cost']
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'IRK', 'integrationOptions': {'numApproxSteps': 1},'energyOptimal': True}
+        solver = casadiSolver(train, track, opts)
+        df, stats = solver.solve(duration)
+
+        energy_IRK_Approx = stats['Cost']
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'IRK', 'integrationOptions': {'numApproxSteps': 0},'energyOptimal': True}
+        solver = casadiSolver(train, track, opts)
+        df, stats = solver.solve(duration)
+
+        energy_IRK = stats['Cost']
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'CVODES', 'energyOptimal': True}
+        solver = casadiSolver(train, track, opts)
+        df, stats = solver.solve(duration)
+
+        energy_CVODES = stats['Cost']
+
+        relDiff_RKApprox_IRKApprox = abs(energy_RK_Approx - energy_IRK_Approx) / energy_RK_Approx
+        relDiff_RKApprox_CVODES = abs(energy_RK_Approx - energy_CVODES) / energy_RK_Approx
+        relDiff_RK_IRK = abs(energy_RK - energy_IRK) / energy_RK
+
+        self.assertLess(
+            relDiff_RKApprox_IRKApprox,
+            tol,
+            msg=(
+                "RK and IRK with numApproxSteps=1 should give similar costs. "
+                f"RK approx: {energy_RK_Approx:.6f}, "
+                f"IRK approx: {energy_IRK_Approx:.6f}, "
+                f"relative difference: {relDiff_RKApprox_IRKApprox:.6f}."
+            )
+        )
+
+        self.assertLess(
+            relDiff_RKApprox_CVODES,
+            tol,
+            msg=(
+                "RK with numApproxSteps=1 and CVODES should give similar costs. "
+                f"RK approx: {energy_RK_Approx:.6f}, "
+                f"CVODES: {energy_CVODES:.6f}, "
+                f"relative difference: {relDiff_RKApprox_CVODES:.6f}."
+            )
+        )
+
+        self.assertLess(
+            relDiff_RK_IRK,
+            tol,
+            msg=(
+                "RK and IRK with numApproxSteps=0 should give similar costs. "
+                f"RK: {energy_RK:.6f}, "
+                f"IRK: {energy_IRK:.6f}, "
+                f"relative difference: {relDiff_RK_IRK:.6f}."
+            )
+        )
\ No newline at end of file

From 6aab5442281c3d0660d38c0502d50e4384cba59d Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Wed, 27 May 2026 16:33:55 +0200
Subject: [PATCH 18/31] implementation of pull request feedback

---
 track.py | 145 +++++++++++++++++++++++++++++++------------------------
 train.py |  28 +++++------
 2 files changed, 97 insertions(+), 76 deletions(-)

diff --git a/track.py b/track.py
index 23392c8..9b0f32e 100644
--- a/track.py
+++ b/track.py
@@ -302,7 +302,7 @@ def checkFields(self):
 
         if not self.crossSectionsOk():
 
-            raise ValueError("Issue with track cross sections!")
+            raise ValueError("Issue with tunnel cross sections!")
 
 
     def importGradientTuples(self, tuples, unit='permil'):
@@ -448,12 +448,12 @@ def importTunnelTuples(self, tuples, unitLength='m', unitCrossSection='m^2'):
             raise ValueError("Specified tunnel units not supported!")
 
         tuples = [(p, convertUnit(l, unitLength), convertUnit(c, unitCrossSection)) for p,l,c in tuples]
-        self.tunnels = importTuples(tuples, 'Position [m]', ['Length [m]', 'CrossSection [m^2]'])
+        self.crossSections = importTuples(tuples, 'Position [m]', ['Length [m]', 'CrossSection [m^2]'])
 
 
         # get end of tunnel positions and assign them a cross section of inf
-        positions = self.tunnels.index.astype(float)
-        tunnelLengths = self.tunnels["Length [m]"].astype(float)
+        positions = self.crossSections.index.astype(float)
+        tunnelLengths = self.crossSections["Length [m]"].astype(float)
 
         endOfTunnelPositions = positions + tunnelLengths
 
@@ -463,19 +463,16 @@ def importTunnelTuples(self, tuples, unitLength='m', unitCrossSection='m^2'):
             if not any(abs((p - e)) < 0.1 for p in positions)
         ]
 
-        openTrack_df = pd.DataFrame({"Length [m]": 0.0, "CrossSection [m^2]": float("inf")}, index=endOfTunnelPositions)
-        openTrack_df.index.name = self.tunnels.index.name
-        self.tunnels = pd.concat([self.tunnels, openTrack_df]).sort_index()
+        nonTunnelSections_df = pd.DataFrame({"Length [m]": 0.0, "CrossSection [m^2]": float("inf")}, index=endOfTunnelPositions)
+        nonTunnelSections_df.index.name = self.crossSections.index.name
+        self.crossSections = pd.concat([self.crossSections, nonTunnelSections_df]).sort_index()
 
         if positions[0] != 0.0:
             first_row = {"Position [m]": 0.0, "Length [m]": 0.0, "CrossSection [m^2]": float("inf")}
-            self.tunnels.loc[0] = first_row
-            self.tunnels = self.tunnels.sort_index()
+            self.crossSections.loc[0] = first_row
+            self.crossSections = self.crossSections.sort_index()
 
-        self.tunnels.drop(columns=["Length [m]"], inplace=True)
-
-        self.crossSections = self.tunnels
-        del self.tunnels
+        self.crossSections.drop(columns=["Length [m]"], inplace=True)
 
         checkDataFrame(self.crossSections, self.length)
 
@@ -630,88 +627,112 @@ def updateSpeedLimitsToTrainLength(self, trainLength):
 
 
     def updateGradientsToTrainLength(self, trainLength):
+        """
+        Convert pointwise gradient changes into train-length-dependent piecewise linear gradients.
 
-        g = self.gradients["Gradient [permil]"].to_numpy(dtype=float)
-        pos = self.gradients.index.to_numpy(dtype=float)
-        slopes = np.r_[0,(g[1:]-g[:-1])/trainLength]
+        A gradient step at position s does not affect the whole train at once.
+        Instead, its effect is spread over one train length, from s to s + trainLength.
+        """
 
-        if len(pos) > 1:
+        gradientValues = self.gradients["Gradient [permil]"].to_numpy(dtype=float)
+        gradientPositions = self.gradients.index.to_numpy(dtype=float)
 
-            pos_adj = np.sort(np.unique(np.r_[pos, pos + trainLength]))
+        if len(gradientPositions) <= 1:
 
-            pos_adj = pos_adj[pos_adj < self.length]
+            return
 
-            g_adj = [g[0]]
-            g_linear = [0]
+        # Each original gradient jump is spread linearly over one train length assuming uniform mass.
+        # The first point has no previous gradient, so its slope contribution is zero.
+        gradientJumpSlopes = np.r_[0.0, (gradientValues[1:] - gradientValues[:-1]) / trainLength]
 
-            for idx in range(1,len(pos_adj)):
+        # New breakpoints occur both when the front of the train reaches a gradient and when the rear of the train has passed it.
+        adjustedPositions = np.sort(np.unique(np.r_[gradientPositions, gradientPositions + trainLength]))
+        adjustedPositions = adjustedPositions[adjustedPositions < self.length]
 
-                currentPosition = pos_adj[idx]
-                previousPosition = pos_adj[idx - 1]
+        adjustedGradients = [gradientValues[0]]
+        gradientLinearTerms = [0.0]
 
-                currentGradient = g_adj[idx-1] + (currentPosition-previousPosition)*g_linear[idx-1]
+        epsilon = 1e-3
 
-                epsilon = 0.001
-                list_indices = []
+        for idx in range(1,len(adjustedPositions)):
 
-                for idx2 in range(len(pos)):
+            currentPosition = adjustedPositions[idx]
+            previousPosition = adjustedPositions[idx - 1]
 
-                    if currentPosition - trainLength + epsilon < pos[idx2] < currentPosition + epsilon:
-                        list_indices.append(idx2)
+            intervalLength = currentPosition - previousPosition
 
-                currentLinearTerm = sum(slopes[list_indices])
+            # Continue the previous linear gradient to the current position.
+            currentGradient = adjustedGradients[-1] + intervalLength * gradientLinearTerms[-1]
 
-                g_adj.append(currentGradient)
-                g_linear.append(currentLinearTerm)
+            # Active gradient jumps are those currently within one train length behind the train front.
+            activeJumpMask = (
+                    (currentPosition - trainLength + epsilon < gradientPositions)
+                    & (gradientPositions < currentPosition + epsilon)
+            )
 
-            # plotGradients(self, np.asarray(pos_adj, dtype=float), np.asarray(g_adj, dtype=float), np.asarray(g_linear, dtype=float))
+            currentLinearTerm = np.sum(gradientJumpSlopes[activeJumpMask])
+
+            adjustedGradients.append(currentGradient)
+            gradientLinearTerms.append(currentLinearTerm)
+
+        # plotGradients(self, np.asarray(pos_adj, dtype=float), np.asarray(g_adj, dtype=float), np.asarray(g_linear, dtype=float))
+
+        self.gradients = pd.DataFrame({"Gradient [permil]": adjustedGradients, "Gradient linear term [permil/m]": gradientLinearTerms}, index=adjustedPositions)
 
-            self.gradients = pd.DataFrame(
-                {"Gradient [permil]": g_adj, "Gradient linear term [permil/m]": g_linear},
-                index=pos_adj,
-            )
 
     def updateCurvaturesToTrainLength(self, trainLength):
+        """
+        Convert pointwise curvature changes into train-length-dependent piecewise linear curvatures.
 
-        c = self.curvatures["Curvature [1/m]"].to_numpy(dtype=float)
-        pos = self.curvatures.index.to_numpy(dtype=float)
-        slopes = np.r_[0, (c[1:] - c[:-1]) / trainLength]
+        A curvature step at position s does not affect the whole train at once.
+        Instead, its effect is spread over one train length, from s to s + trainLength.
+        """
 
-        if len(pos) > 1:
+        curvatureValues = self.curvatures["Curvature [1/m]"].to_numpy(dtype=float)
+        curvaturePositions = self.curvatures.index.to_numpy(dtype=float)
 
-            pos_adj = np.sort(np.r_[pos, pos + trainLength])
+        if len(curvaturePositions) <= 1:
 
-            pos_adj = pos_adj[pos_adj < self.length]
+            return
 
-            c_adj = [c[0]]
-            c_linear = [0]
+        # Each original curvature jump is spread linearly over one train length assuming uniform mass.
+        # The first point has no previous curvature, so its slope contribution is zero.
+        curvatureJumpSlopes = np.r_[0.0, (curvatureValues[1:] - curvatureValues[:-1]) / trainLength]
 
-            for idx in range(1,len(pos_adj)):
+        # New breakpoints occur both when the front of the train reaches a curvature
+        # change and when the rear of the train has passed it.
+        adjustedPositions = np.sort(np.unique(np.r_[curvaturePositions, curvaturePositions + trainLength]))
+        adjustedPositions = adjustedPositions[adjustedPositions < self.length]
 
-                currentPosition = pos_adj[idx]
-                previousPosition = pos_adj[idx - 1]
+        adjustedCurvatures = [curvatureValues[0]]
+        curvatureLinearTerms = [0.0]
 
-                currentCurvature = c_adj[idx-1] + (currentPosition-previousPosition)*c_linear[idx-1]
+        epsilon = 1e-3
 
-                epsilon = 0.001
-                list_indices = []
+        for idx in range(1, len(adjustedPositions)):
 
-                for idx2 in range(len(pos)):
+            currentPosition = adjustedPositions[idx]
+            previousPosition = adjustedPositions[idx - 1]
 
-                    if currentPosition - trainLength + epsilon < pos[idx2] < currentPosition + epsilon:
-                        list_indices.append(idx2)
+            intervalLength = currentPosition - previousPosition
 
-                currentLinearTerm = sum(slopes[list_indices])
+            # Continue the previous linear curvature to the current position.
+            currentCurvature = adjustedCurvatures[-1] + intervalLength * curvatureLinearTerms[-1]
 
-                c_adj.append(currentCurvature)
-                c_linear.append(currentLinearTerm)
+            # Active curvature jumps are those currently within one train length behind the train front.
+            activeJumpMask = (
+                    (currentPosition - trainLength + epsilon < curvaturePositions)
+                    & (curvaturePositions < currentPosition + epsilon)
+            )
+
+            currentLinearTerm = np.sum(curvatureJumpSlopes[activeJumpMask])
+
+            adjustedCurvatures.append(currentCurvature)
+            curvatureLinearTerms.append(currentLinearTerm)
 
             # plotCurvatures(self, np.asarray(pos_adj, dtype=float), np.asarray(c_adj, dtype=float), np.asarray(c_linear, dtype=float))
 
-            self.curvatures = pd.DataFrame(
-                {"Curvature [1/m]": c_adj, "Curvature linear term [1/m^2]": c_linear},
-                index=pos_adj,
-            )
+            self.curvatures = pd.DataFrame({"Curvature [1/m]": adjustedCurvatures, "Curvature linear term [1/m^2]": curvatureLinearTerms}, index=adjustedPositions)
 
 
 if __name__ == '__main__':
diff --git a/train.py b/train.py
index fbfb2c7..1ee7100 100644
--- a/train.py
+++ b/train.py
@@ -246,38 +246,38 @@ def __init__(self, sr0, sr1, sr2, rho=1, g=9.81, withPnBrake=True) -> None:
 
         # states
 
-        time = ca.MX.sym('time')                            # [s]
-        velocitySquared = ca.MX.sym('velocitySquared')      # [m^2/s^2]
-        position = ca.MX.sym('position')                    # [m]
+        time = ca.MX.sym('time')  # [s]
+        velocitySquared = ca.MX.sym('velocitySquared')  # [m^2/s^2]
+        position = ca.MX.sym('position')  # [m]
 
         x = ca.vertcat(time, velocitySquared, position)
 
         # controls
 
-        traction = ca.MX.sym('traction')                    # [N/kg]
-        pnBrake = ca.MX.sym('pnBrake')                      # [N/kg]
+        traction = ca.MX.sym('traction')  # [N/kg]
+        pnBrake = ca.MX.sym('pnBrake') # [N/kg]
 
         u = ca.vertcat(traction, pnBrake if withPnBrake else [])
 
         # parameters
 
-        gradient = ca.MX.sym('gradient')                        # [-]     -> values between [0,0.2]
-        gradientLinearTerm = ca.MX.sym('gradientLinearTerm')    # [1/m]
-        curvature = ca.MX.sym('curvature')                      # [1/m]   -> values between [0,0.004]
+        gradient = ca.MX.sym('gradient')  # [-]
+        gradientLinearTerm = ca.MX.sym('gradientLinearTerm')  # [1/m]
+        curvature = ca.MX.sym('curvature')  # [1/m]
         curvatureLinearTerm = ca.MX.sym('curvatureLinearTerm')  # [1/m^2]
-        tunnelFactor = ca.MX.sym('tunnelFactor')                # [1/m]
+        tunnelFactor = ca.MX.sym('tunnelFactor')  # [1/m]
         ds = ca.MX.sym('ds')
 
         p = ca.vertcat(gradient, gradientLinearTerm, curvature, curvatureLinearTerm, tunnelFactor, ds)
 
         # ODE
 
-        rollingResistance = sr0 + sr1*ca.sqrt(velocitySquared) + sr2*velocitySquared # [N/kg]
-        gradientResistance = g*(1/rho)*(gradient+gradientLinearTerm*position) # [N/kg]
-        curvatureResistance = (1/rho)*(5.07 * (curvature+curvatureLinearTerm*position)) # [N/kg]
-        tunnelResistance = tunnelFactor * velocitySquared # [N/kg]
+        rollingResistance = sr0 + sr1*ca.sqrt(velocitySquared) + sr2*velocitySquared  # [N/kg]
+        gradientResistance = g*(1/rho)*(gradient+gradientLinearTerm*position)  # [N/kg]
+        curvatureResistance = (1/rho)*(5.07 * (curvature+curvatureLinearTerm*position))  # [N/kg]
+        tunnelResistance = tunnelFactor * velocitySquared  # [N/kg]
 
-        acceleration = traction + (pnBrake if withPnBrake else 0) - rollingResistance - gradientResistance - curvatureResistance - tunnelResistance # [m/s^2]
+        acceleration = traction + (pnBrake if withPnBrake else 0) - rollingResistance - gradientResistance - curvatureResistance - tunnelResistance  # [m/s^2]
 
         timeODE = 1/ca.sqrt(velocitySquared)
         velocityODE = 2*acceleration

From 32d3c76b7e7b84759b84e77f5aa727031911d05c Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Wed, 27 May 2026 16:41:48 +0200
Subject: [PATCH 19/31] bug fixes

---
 track.py                                                    | 4 ++--
 unitTests/trainLengthDependentTrackAttributes/speedLimit.py | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/track.py b/track.py
index 9b0f32e..4fcfed1 100644
--- a/track.py
+++ b/track.py
@@ -730,9 +730,9 @@ def updateCurvaturesToTrainLength(self, trainLength):
             adjustedCurvatures.append(currentCurvature)
             curvatureLinearTerms.append(currentLinearTerm)
 
-            # plotCurvatures(self, np.asarray(pos_adj, dtype=float), np.asarray(c_adj, dtype=float), np.asarray(c_linear, dtype=float))
+        # plotCurvatures(self, np.asarray(pos_adj, dtype=float), np.asarray(c_adj, dtype=float), np.asarray(c_linear, dtype=float))
 
-            self.curvatures = pd.DataFrame({"Curvature [1/m]": adjustedCurvatures, "Curvature linear term [1/m^2]": curvatureLinearTerms}, index=adjustedPositions)
+        self.curvatures = pd.DataFrame({"Curvature [1/m]": adjustedCurvatures, "Curvature linear term [1/m^2]": curvatureLinearTerms}, index=adjustedPositions)
 
 
 if __name__ == '__main__':
diff --git a/unitTests/trainLengthDependentTrackAttributes/speedLimit.py b/unitTests/trainLengthDependentTrackAttributes/speedLimit.py
index 68fba47..d5a9a1a 100644
--- a/unitTests/trainLengthDependentTrackAttributes/speedLimit.py
+++ b/unitTests/trainLengthDependentTrackAttributes/speedLimit.py
@@ -22,7 +22,7 @@ def testTrainDoesNotAccelerateToEarly(self):
 
         train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='trains')
 
-        track = Track(config={'id': 'test_speed_increase'}, pathJSON='tracks')
+        track = Track(config={'id': 'test_one_speed_increase'}, pathJSON='tracks')
         track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
 
         opts = {'numIntervals': 300, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True}

From 281cb1424d5f08bc561dbdaf1fd5288e9cc31458 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Mon, 1 Jun 2026 15:40:31 +0200
Subject: [PATCH 20/31] wip

---
 ocp.py                                        |  13 +-
 simulations/sim_launcher.py                   |   4 +-
 tracks/swisstopo/analyzeTracks.py             | 113 ++++++++
 tracks/swisstopo/csvImporter.py               | 108 ++++++++
 tracks/test_two_radii.json                    |  12 +-
 train.py                                      |   8 +-
 ...irt_Tpf.json => CH_Stadler_FLIRT_TPF.json} |   0
 unitTests/integrators/integrators.py          |  99 +++++++
 .../gradient.py                               | 242 +++++++++++-------
 9 files changed, 488 insertions(+), 111 deletions(-)
 create mode 100644 tracks/swisstopo/analyzeTracks.py
 create mode 100644 tracks/swisstopo/csvImporter.py
 rename trains/{Flirt_Tpf.json => CH_Stadler_FLIRT_TPF.json} (100%)
 create mode 100644 unitTests/integrators/integrators.py

diff --git a/ocp.py b/ocp.py
index b76cdd6..68fe534 100644
--- a/ocp.py
+++ b/ocp.py
@@ -209,9 +209,16 @@ def __init__(self, train, track, optsDict={}):
                 ubg += [accMax]
 
                 # coupling constraints
-                out = trainIntegrator.solve(time=time[i], velocitySquared=velSq[i], ds=self.steps[i],
-                    traction=Fel[i], pnBrake=Fpb[i], gradient=grad, gradientLinearTerm=curvLinearTerm, curvature=curv,
-                                            curvatureLinearTerm=curvLinearTerm, tunnelFactor=tunnelFactor)
+                out = trainIntegrator.solve(time=time[i],
+                                            velocitySquared=velSq[i],
+                                            ds=self.steps[i],
+                                            traction=Fel[i],
+                                            pnBrake=Fpb[i],
+                                            gradient=grad,
+                                            gradientLinearTerm=gradLinearTerm,
+                                            curvature=curv,
+                                            curvatureLinearTerm=curvLinearTerm,
+                                            tunnelFactor=tunnelFactor)
 
                 xNxt1 = ca.vertcat(time[i+1], velSq[i+1])
                 xNxt2 = ca.vertcat(out['time'], out['velSquared'])
diff --git a/simulations/sim_launcher.py b/simulations/sim_launcher.py
index 558243c..93b70d6 100644
--- a/simulations/sim_launcher.py
+++ b/simulations/sim_launcher.py
@@ -13,11 +13,11 @@
     train = Train(config={'id':'Flirt_Tpf'}, pathJSON='../trains')
 
     track = Track(config={'id':'CH_ZH_LU'}, pathJSON='../tracks')
-    # track = Track(config={'id':'CH_StGallen_Wil'}, pathJSON='../tracks')
+    track = Track(config={'id':'CH_StGallen_Wil'}, pathJSON='../tracks')
     track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
     track.updateTrainLengthDependentValues(train)
 
-    opts = {'numIntervals':600, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':0}, 'energyOptimal':True}
+    opts = {'numIntervals':600, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':1}, 'energyOptimal':True}
 
     solver = casadiSolver(train, track, opts)
 
diff --git a/tracks/swisstopo/analyzeTracks.py b/tracks/swisstopo/analyzeTracks.py
new file mode 100644
index 0000000..4f69d9c
--- /dev/null
+++ b/tracks/swisstopo/analyzeTracks.py
@@ -0,0 +1,113 @@
+import numpy as np
+from matplotlib import pyplot as plt
+
+from ocp import casadiSolver
+from track import Track
+from train import Train
+
+
+if __name__ == '__main__':
+
+
+    SBB_track = Track(config={'id': 'CH_StGallen_Wil'}, pathJSON='')
+
+    SBB_positions = SBB_track.gradients.index.values
+    SBB_gradients = SBB_track.gradients["Gradient [permil]"].to_numpy()
+
+    initial_altitude = SBB_track.altitude
+    delta_s = np.diff(SBB_positions)
+    delta_h = SBB_gradients[:-1] / 1000 * delta_s
+
+    SBB_altitude = np.insert(initial_altitude + np.cumsum(delta_h),0, initial_altitude)
+
+    Topo_track = Track(config={'id': 'CH_StGallen_Wil_Swisstopo'}, pathJSON='')
+
+    Topo_positions = Topo_track.gradients.index.values
+    Topo_gradients = Topo_track.gradients["Gradient [permil]"].to_numpy()
+
+    initial_altitude = Topo_track.altitude
+    delta_s = np.diff(Topo_positions)
+    delta_h = Topo_gradients[:-1] / 1000 * delta_s
+
+    Topo_altitude = np.insert(initial_altitude + np.cumsum(delta_h),0, initial_altitude)
+
+    shift = 800 # 770
+
+
+    ### Plot Altitude Comparison
+
+    fig, ax = plt.subplots(figsize=(16, 8))
+
+    ax.plot(SBB_positions / 1000, SBB_altitude, label="SBB")
+    ax.plot((Topo_positions - shift) / 1000, Topo_altitude, label="Topo")
+    ax.set_title("Altitude Comparison")
+    ax.set_xlabel("Position [km]")
+    ax.set_ylabel("Altitude [m]")
+    ax.grid(True, which="both", linestyle="--", alpha=0.5)
+    ax.legend(loc="upper right")
+    ax.set_xlim(0, SBB_track.length / 1000)
+    ax.figure.tight_layout()
+
+    plt.show()
+
+
+    ### Plot Speed Limit Comparison
+
+    fig2, ax2 = plt.subplots(figsize=(16, 8))
+
+    ax2.step(SBB_track.speedLimits.index.values / 1000, SBB_track.speedLimits["Speed limit [m/s]"].to_numpy()*3.6, label="SBB")
+    ax2.step((Topo_track.speedLimits.index.values-shift) / 1000, Topo_track.speedLimits["Speed limit [m/s]"].to_numpy()*3.6, label="Topo")
+    ax2.set_title("Altitude Comparison")
+    ax2.set_xlabel("Position [km]")
+    ax2.set_ylabel("Velocity [km/h]")
+    ax2.grid(True, which="both", linestyle="--", alpha=0.5)
+    ax2.legend(loc="upper right")
+    ax2.set_xlim(0, SBB_track.length / 1000)
+    ax2.figure.tight_layout()
+
+    plt.show()
+
+
+    ### Compute Energy Comparison
+
+    # Timetable
+    startPosition = 0           # [m]
+    endPosition = 23000         # [m]
+    duration = 23000/(80/3.6)   # [s]
+
+    train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='../../trains')
+    opts = {'numIntervals': 1000, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 0}, 'energyOptimal': True}
+
+    SBB_track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
+    SBB_track.updateTrainLengthDependentValues(train)
+    solver = casadiSolver(train, SBB_track, opts)
+    dfSBB, statsSBB = solver.solve(duration)
+
+    Topo_track.updateLimits(positionStart=startPosition + shift, positionEnd=endPosition + shift, unit='m')
+    Topo_track.updateTrainLengthDependentValues(train)
+    solver = casadiSolver(train, Topo_track, opts)
+    dfTopo, statsTopo = solver.solve(duration)
+
+    print(f"Cost SBB: {statsSBB['Cost']:.2f}")
+    print(f"Cost Topo: {statsTopo['Cost']:.2f}")
+
+    print(f"{abs(statsSBB['Cost'] - statsTopo['Cost']) / statsSBB['Cost'] * 100:.2f}%")
+
+
+    ### Plot Trajectory
+
+    fig3, ax3 = plt.subplots(figsize=(16, 8))
+
+    ax3.plot(dfSBB["Position [m]"] / 1000, dfSBB["Velocity [m/s]"] * 3.6, label="SBB")
+    ax3.plot(dfTopo["Position [m]"] / 1000, dfTopo["Velocity [m/s]"] * 3.6, label="Topo")
+    ax3.set_title("Speed Profile Comparison")
+    ax3.set_xlabel("Position [km]")
+    ax3.set_ylabel("Velocity [km/h]")
+    ax3.grid(True, which="both", linestyle="--", alpha=0.5)
+    ax3.legend(loc="upper right")
+    ax3.set_xlim(0, dfSBB["Position [m]"].max() / 1000)
+    ax3.figure.tight_layout()
+
+    plt.show()
+
+
diff --git a/tracks/swisstopo/csvImporter.py b/tracks/swisstopo/csvImporter.py
new file mode 100644
index 0000000..f448c28
--- /dev/null
+++ b/tracks/swisstopo/csvImporter.py
@@ -0,0 +1,108 @@
+import json
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+if __name__ == '__main__':
+
+
+    ### Read CSV
+
+    filePath = r"C:\Users\rolan\Documents\ms-eetc-innocheque\tracks\swisstopo\Track_StGallen_Wil.csv"
+
+    df = pd.read_csv(filePath,na_values=["<null>", "null", ""])
+
+    print(df.head())
+    print(df.dtypes)
+
+
+    ### Speed limits
+
+    speedProfile = df.loc[
+        df["V_max"].notna(),
+        ["Total_Distance", "V_max"]
+    ].copy()
+
+
+    ### Gradients
+
+    totalDistance = df["Total_Distance"]
+    altitude = df["Altitude"]
+
+    window_size = 7
+    altitude = altitude.rolling(window=window_size, center=True, min_periods=1).mean()
+
+    spacing = 50  # [m]
+
+    positions = np.arange(0, totalDistance.max(), spacing)
+    altitude_interp = np.interp(positions, totalDistance, altitude)
+
+    gradientPerMille = np.insert(1000 * np.diff(altitude_interp) / np.diff(positions),0,0 )
+    gradientPerMille = np.round(gradientPerMille, 1)
+
+
+    ### Parse to Json
+
+    track_id = "CH_StGallen_Wil_Swisstopo"
+    author = "Roland Staerk"
+
+    name = "CH_StGallen_Wil_Swisstopo"
+
+    output_dir = Path(r"C:\Users\rolan\Documents\ms-eetc-innocheque\tracks\swisstopo")
+    output_path = output_dir / f"{name}.json"
+
+    stops = [
+        0.0,
+        float(totalDistance.iloc[-1])
+    ]
+
+    speed_limits = [
+        [float(pos), float(vmax)]
+        for pos, vmax in zip(
+            speedProfile["Total_Distance"],
+            speedProfile["V_max"]
+        )
+    ]
+
+    gradients = [
+        [float(pos), float(grad)]
+        for pos, grad in zip(positions, gradientPerMille)
+    ]
+
+    track_data = {
+        "metadata": {
+            "id": track_id,
+            "created by": author,
+            "library version": "TTOBench v1.4",
+            "license": "BSD 2-Clause License"
+        },
+        "altitude": {
+            "unit": "m",
+            "value": float(altitude_interp[0])
+        },
+        "stops": {
+            "unit": "m",
+            "values": stops
+        },
+        "speed limits": {
+            "units": {
+                "position": "m",
+                "velocity": "km/h"
+            },
+            "values": speed_limits
+        },
+        "gradients": {
+            "units": {
+                "position": "m",
+                "slope": "permil"
+            },
+            "values": gradients
+        }
+    }
+
+
+    ### Save Json
+
+    with open(output_path, "w", encoding="utf-8") as f:
+        json.dump(track_data, f, indent=4)
\ No newline at end of file
diff --git a/tracks/test_two_radii.json b/tracks/test_two_radii.json
index 12b67c0..27fd8b7 100644
--- a/tracks/test_two_radii.json
+++ b/tracks/test_two_radii.json
@@ -37,8 +37,8 @@
         "values": [
             [
                 0.0,
-                Infinity,
-                Infinity
+                "infinity",
+                "infinity"
             ],
             [
                 1000.0,
@@ -47,8 +47,8 @@
             ],
             [
                 2000.0,
-                Infinity,
-                Infinity
+                "infinity",
+                "infinity"
             ],
             [
                 3000.0,
@@ -57,8 +57,8 @@
             ],
             [
                 4000.0,
-                Infinity,
-                Infinity
+                "infinity",
+                "infinity"
             ]
         ]
     }
diff --git a/train.py b/train.py
index 1ee7100..bb62eae 100644
--- a/train.py
+++ b/train.py
@@ -372,8 +372,12 @@ def __init__(self, model, solver, optsDict={}) -> None:
 
             epsVelSq = 0.0001
             for idx in range(ns):
-                vCurr = ca.sqrt(ca.fmax(zf[0, idx], epsVelSq))
-                vNext = ca.sqrt(ca.fmax(zf[0, idx + 1], epsVelSq))
+                # vCurr = ca.sqrt(ca.fmax(zf[0, idx], epsVelSq))
+                # vNext = ca.sqrt(ca.fmax(zf[0, idx + 1], epsVelSq))
+
+                vCurr = ca.sqrt(zf[0, idx])
+                vNext = ca.sqrt(zf[0, idx + 1])
+
                 tApprox += 2*model.parameters[-1]*(evalPoints[idx+1]-evalPoints[idx])/(vCurr + vNext)
 
             eval = ca.vertcat(tApprox, zf[0, ns], zf[1, ns])
diff --git a/trains/Flirt_Tpf.json b/trains/CH_Stadler_FLIRT_TPF.json
similarity index 100%
rename from trains/Flirt_Tpf.json
rename to trains/CH_Stadler_FLIRT_TPF.json
diff --git a/unitTests/integrators/integrators.py b/unitTests/integrators/integrators.py
new file mode 100644
index 0000000..f6fa23d
--- /dev/null
+++ b/unitTests/integrators/integrators.py
@@ -0,0 +1,99 @@
+import unittest
+
+from ocp import casadiSolver
+from track import Track
+from train import Train
+
+
+class TestGradient(unittest.TestCase):
+
+    def testAllIntegratorTypesWork(self):
+        '''
+        Verify that all supported integration methods produce consistent results
+        for the same train, track, and optimization setup.
+
+        The test compares RK, IRK, and CVODES, including the approximate time
+        integration option for RK and IRK. The resulting energy costs should only
+        differ by a small relative tolerance.
+        '''
+
+        startPosition = 0  # [m]
+        endPosition = 5000  # [m]
+        duration = 5000 / (60 / 3.6)  # [s]
+
+        tol = 0.02
+        numIntervals = 100
+
+        train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='trains')
+        train.length = 600
+
+        track = Track(config={'id': 'test_one_hill'}, pathJSON='tracks')
+        track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
+        track.updateTrainLengthDependentValues(train)
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}}
+        solver = casadiSolver(train, track, opts)
+        df, stats = solver.solve(duration)
+
+        energy_RK_Approx = stats['Cost']
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 0}}
+        solver = casadiSolver(train, track, opts)
+        df, stats = solver.solve(duration)
+
+        energy_RK = stats['Cost']
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'IRK', 'integrationOptions': {'numApproxSteps': 1}}
+        solver = casadiSolver(train, track, opts)
+        df, stats = solver.solve(duration)
+
+        energy_IRK_Approx = stats['Cost']
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'IRK', 'integrationOptions': {'numApproxSteps': 0}}
+        solver = casadiSolver(train, track, opts)
+        df, stats = solver.solve(duration)
+
+        energy_IRK = stats['Cost']
+
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'CVODES'}
+        solver = casadiSolver(train, track, opts)
+        df, stats = solver.solve(duration)
+
+        energy_CVODES = stats['Cost']
+
+        relDiff_RKApprox_IRKApprox = abs(energy_RK_Approx - energy_IRK_Approx) / energy_RK_Approx
+        relDiff_RKApprox_CVODES = abs(energy_RK_Approx - energy_CVODES) / energy_RK_Approx
+        relDiff_RK_IRK = abs(energy_RK - energy_IRK) / energy_RK
+
+        self.assertLess(
+            relDiff_RKApprox_IRKApprox,
+            tol,
+            msg=(
+                "RK and IRK with numApproxSteps=1 should give similar costs. "
+                f"RK approx: {energy_RK_Approx:.6f}, "
+                f"IRK approx: {energy_IRK_Approx:.6f}, "
+                f"relative difference: {relDiff_RKApprox_IRKApprox:.6f}."
+            )
+        )
+
+        self.assertLess(
+            relDiff_RKApprox_CVODES,
+            tol,
+            msg=(
+                "RK with numApproxSteps=1 and CVODES should give similar costs. "
+                f"RK approx: {energy_RK_Approx:.6f}, "
+                f"CVODES: {energy_CVODES:.6f}, "
+                f"relative difference: {relDiff_RKApprox_CVODES:.6f}."
+            )
+        )
+
+        self.assertLess(
+            relDiff_RK_IRK,
+            tol,
+            msg=(
+                "RK and IRK with numApproxSteps=0 should give similar costs. "
+                f"RK: {energy_RK:.6f}, "
+                f"IRK: {energy_IRK:.6f}, "
+                f"relative difference: {relDiff_RK_IRK:.6f}."
+            )
+        )
\ No newline at end of file
diff --git a/unitTests/trainLengthDependentTrackAttributes/gradient.py b/unitTests/trainLengthDependentTrackAttributes/gradient.py
index 7b089aa..e5a9d4a 100644
--- a/unitTests/trainLengthDependentTrackAttributes/gradient.py
+++ b/unitTests/trainLengthDependentTrackAttributes/gradient.py
@@ -1,14 +1,151 @@
 import unittest
 
+import numpy as np
 from matplotlib import pyplot as plt
 
-from ocp import casadiSolver
+from ocp import casadiSolver, OptionsCasadiSolver
 from track import Track, computeAltitude
-from train import Train
+from train import Train, TrainIntegrator
 
 
 class TestGradient(unittest.TestCase):
 
+    def test_integrator_CVODES(self):
+        '''
+        Track with a linearly increasing gradient over 1000 m.
+
+        The result obtained using the piecewise linear gradient model is compared
+        against a piecewise constant midpoint approximation with increasing numbers
+        of intervals.
+
+        The piecewise constant approximation should converge to the piecewise linear
+        result for both duration and final velocity.
+        '''
+
+        train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
+
+        optsDict = {'integrationMethod': 'CVODES'}
+        opts = OptionsCasadiSolver(optsDict)
+
+        trainModel = train.exportModel()
+        trainIntegrator = TrainIntegrator(trainModel, opts.integrationMethod, opts.integrationOptions.toDict())
+
+        # Scenario
+        time0 = 0
+        velSq0 = 1
+        ds = 1000
+        traction = 0.8 * (train.forceMax / train.mass)
+
+        initialGradient = 0
+        finalGradient = 0.07
+
+        maxIntervals = 50
+        relativeTolerance = 1e-3
+        plotDebug = True
+
+        # PWL gradient reference
+        out = trainIntegrator.solve(
+            time=time0,
+            velocitySquared=velSq0,
+            ds=ds,
+            traction=traction,
+            pnBrake=0,
+            gradient=initialGradient,
+            gradientLinearTerm=(finalGradient - initialGradient) / ds,
+            curvature=0,
+            curvatureLinearTerm=0,
+            tunnelFactor=0
+        )
+
+        pwlDuration = float(out['time'])
+        pwlVelocity = np.sqrt(float(out['velSquared']))
+
+        # PWC gradients using midpoint rule
+        times = []
+        velocities = []
+        intervalCounts = []
+
+        for numIntervals in range(1, maxIntervals + 1):
+
+            time = time0
+            velSq = velSq0
+
+            for idx in range(numIntervals):
+
+                gradient = (initialGradient + (idx + 0.5) * (finalGradient - initialGradient) / numIntervals)
+
+                out = trainIntegrator.solve(
+                    time=time,
+                    velocitySquared=velSq,
+                    ds=ds / numIntervals,
+                    traction=traction,
+                    pnBrake=0,
+                    gradient=gradient,
+                    gradientLinearTerm=0,
+                    curvature=0,
+                    curvatureLinearTerm=0,
+                    tunnelFactor=0
+                )
+
+                time = out['time']
+                velSq = out['velSquared']
+
+            intervalCounts.append(numIntervals)
+            times.append(float(time))
+            velocities.append(np.sqrt(float(velSq)))
+
+        finalPwcDuration = times[-1]
+        finalPwcVelocity = velocities[-1]
+
+        relativeDurationError = abs(finalPwcDuration - pwlDuration) / pwlDuration
+        relativeVelocityError = abs(finalPwcVelocity - pwlVelocity) / pwlVelocity
+
+        self.assertLess(
+            relativeDurationError,
+            relativeTolerance,
+            msg=(
+                "PWC midpoint approximation did not converge sufficiently "
+                "to the PWL duration. "
+                f"PWL duration: {pwlDuration:.6f}, "
+                f"PWC duration: {finalPwcDuration:.6f}, "
+                f"relative error: {relativeDurationError:.6e}."
+            )
+        )
+
+        self.assertLess(
+            relativeVelocityError,
+            relativeTolerance,
+            msg=(
+                "PWC midpoint approximation did not converge sufficiently "
+                "to the PWL final velocity. "
+                f"PWL velocity: {pwlVelocity:.6f}, "
+                f"PWC velocity: {finalPwcVelocity:.6f}, "
+                f"relative error: {relativeVelocityError:.6e}."
+            )
+        )
+
+        if plotDebug:
+            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(16, 10))
+
+            ax1.axhline(pwlDuration, label="pwl")
+            ax1.plot(intervalCounts, times, marker="o", color="orange", label="pwc midpoint")
+            ax1.set_xlabel("Number of intervals")
+            ax1.set_ylabel("Duration [s]")
+            ax1.grid(True, which="both", linestyle="--", alpha=0.5)
+            ax1.legend(loc="upper right")
+
+            ax2.axhline(pwlVelocity, label="pwl")
+            ax2.plot(intervalCounts, velocities, marker="o", color="orange", label="pwc midpoint")
+            ax2.set_xlabel("Number of intervals")
+            ax2.set_ylabel("Velocity [m/s]")
+            ax2.grid(True, which="both", linestyle="--", alpha=0.5)
+            ax2.legend(loc="upper right")
+
+            fig.tight_layout()
+            plt.show()
+
+
+
     def testLinearGradient(self):
         '''
         Track with 20 permil increase from 1000 m to 2000 m and 20 permil
@@ -25,23 +162,23 @@ def testLinearGradient(self):
         duration = 5000/(60/3.6) # [s]
 
         altitudeTolerance = 1e-4
-        energyRelativeTolerance = 0.004
+        energyRelativeTolerance = 1e-4
         numIntervals = 100
 
-        train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='trains')
+        train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
         train.length = 600
 
         track = Track(config={'id': 'test_one_hill'}, pathJSON='tracks')
         track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
         track.updateTrainLengthDependentValues(train)
 
-        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True}
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 0}, 'energyOptimal': True}
         solver = casadiSolver(train, track, opts)
         df1, stats1 = solver.solve(duration)
 
         energyConsumptionWithLinearTerms = stats1['Cost']
 
-        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True, 'pwcLengthDependentTrackAttributes': True}
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 0}, 'energyOptimal': True, 'pwcLengthDependentTrackAttributes': True}
         solver = casadiSolver(train, track, opts)
         df2, stats2 = solver.solve(duration)
 
@@ -92,95 +229,4 @@ def testLinearGradient(self):
             ax.set_xlim(0, track.length / 1000)
             ax.figure.tight_layout()
 
-            plt.show()
-
-    def testAllIntegratorTypesWork(self):
-        '''
-        Verify that all supported integration methods produce consistent results
-        for the same train, track, and optimization setup.
-
-        The test compares RK, IRK, and CVODES, including the approximate time
-        integration option for RK and IRK. The resulting energy costs should only
-        differ by a small relative tolerance.
-        '''
-
-        startPosition = 0  # [m]
-        endPosition = 5000  # [m]
-        duration = 5000 / (60 / 3.6)  # [s]
-
-        tol = 0.02
-        numIntervals = 100
-
-        train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='trains')
-        train.length = 600
-
-        track = Track(config={'id': 'test_one_hill'}, pathJSON='tracks')
-        track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
-        track.updateTrainLengthDependentValues(train)
-
-        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1},'energyOptimal': True}
-        solver = casadiSolver(train, track, opts)
-        df, stats = solver.solve(duration)
-
-        energy_RK_Approx = stats['Cost']
-
-        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 0},'energyOptimal': True}
-        solver = casadiSolver(train, track, opts)
-        df, stats = solver.solve(duration)
-
-        energy_RK = stats['Cost']
-
-        opts = {'numIntervals': numIntervals, 'integrationMethod': 'IRK', 'integrationOptions': {'numApproxSteps': 1},'energyOptimal': True}
-        solver = casadiSolver(train, track, opts)
-        df, stats = solver.solve(duration)
-
-        energy_IRK_Approx = stats['Cost']
-
-        opts = {'numIntervals': numIntervals, 'integrationMethod': 'IRK', 'integrationOptions': {'numApproxSteps': 0},'energyOptimal': True}
-        solver = casadiSolver(train, track, opts)
-        df, stats = solver.solve(duration)
-
-        energy_IRK = stats['Cost']
-
-        opts = {'numIntervals': numIntervals, 'integrationMethod': 'CVODES', 'energyOptimal': True}
-        solver = casadiSolver(train, track, opts)
-        df, stats = solver.solve(duration)
-
-        energy_CVODES = stats['Cost']
-
-        relDiff_RKApprox_IRKApprox = abs(energy_RK_Approx - energy_IRK_Approx) / energy_RK_Approx
-        relDiff_RKApprox_CVODES = abs(energy_RK_Approx - energy_CVODES) / energy_RK_Approx
-        relDiff_RK_IRK = abs(energy_RK - energy_IRK) / energy_RK
-
-        self.assertLess(
-            relDiff_RKApprox_IRKApprox,
-            tol,
-            msg=(
-                "RK and IRK with numApproxSteps=1 should give similar costs. "
-                f"RK approx: {energy_RK_Approx:.6f}, "
-                f"IRK approx: {energy_IRK_Approx:.6f}, "
-                f"relative difference: {relDiff_RKApprox_IRKApprox:.6f}."
-            )
-        )
-
-        self.assertLess(
-            relDiff_RKApprox_CVODES,
-            tol,
-            msg=(
-                "RK with numApproxSteps=1 and CVODES should give similar costs. "
-                f"RK approx: {energy_RK_Approx:.6f}, "
-                f"CVODES: {energy_CVODES:.6f}, "
-                f"relative difference: {relDiff_RKApprox_CVODES:.6f}."
-            )
-        )
-
-        self.assertLess(
-            relDiff_RK_IRK,
-            tol,
-            msg=(
-                "RK and IRK with numApproxSteps=0 should give similar costs. "
-                f"RK: {energy_RK:.6f}, "
-                f"IRK: {energy_IRK:.6f}, "
-                f"relative difference: {relDiff_RK_IRK:.6f}."
-            )
-        )
\ No newline at end of file
+            plt.show()
\ No newline at end of file

From 6d4b0af81d3cd77fe5a81a72280c997af7fdb5dc Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Mon, 1 Jun 2026 21:27:10 +0200
Subject: [PATCH 21/31] wip 2

---
 simulations/sim_launcher.py       | 18 ++++++++++++++++++
 tracks/swisstopo/analyzeTracks.py | 10 +++++++---
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/simulations/sim_launcher.py b/simulations/sim_launcher.py
index 93b70d6..9bf3443 100644
--- a/simulations/sim_launcher.py
+++ b/simulations/sim_launcher.py
@@ -1,5 +1,20 @@
+from efficiency import totalLossesFunction
 from ocp import casadiSolver
 
+def get_power_loss_function(train, mode="perfect",* ,auxiliaries: float = 27_000, eta_gear: float = 0.96):
+
+    if mode == "perfect":
+        return lambda f, v: 0
+
+    elif mode == "static":
+        return lambda f, v: (f>0)*f*v*(1-train.etaTraction)/train.etaTraction - (f<0)*f*v*(1-train.etaRgBrake)
+
+    elif mode == "dynamic":
+        return totalLossesFunction(train, auxiliaries=auxiliaries, etaGear=eta_gear)
+
+    else:
+        raise ValueError("mode must be one of: 'perfect', 'static', 'dynamic'")
+
 if __name__ == '__main__':
 
     from train import Train
@@ -11,6 +26,9 @@
     duration = 60*20        # [s]
 
     train = Train(config={'id':'Flirt_Tpf'}, pathJSON='../trains')
+    train.forceMinPn = 0
+    train.withPnBrake = False
+    train.powerLosses = get_power_loss_function(train, "static")
 
     track = Track(config={'id':'CH_ZH_LU'}, pathJSON='../tracks')
     track = Track(config={'id':'CH_StGallen_Wil'}, pathJSON='../tracks')
diff --git a/tracks/swisstopo/analyzeTracks.py b/tracks/swisstopo/analyzeTracks.py
index 4f69d9c..7ce81fe 100644
--- a/tracks/swisstopo/analyzeTracks.py
+++ b/tracks/swisstopo/analyzeTracks.py
@@ -2,6 +2,7 @@
 from matplotlib import pyplot as plt
 
 from ocp import casadiSolver
+from simulations.sim_launcher import get_power_loss_function
 from track import Track
 from train import Train
 
@@ -75,11 +76,14 @@
     endPosition = 23000         # [m]
     duration = 23000/(80/3.6)   # [s]
 
-    train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='../../trains')
-    opts = {'numIntervals': 1000, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 0}, 'energyOptimal': True}
+    train = Train(config={'id': 'CH_Stadler_FLIRT_TPF'}, pathJSON='../../trains')
+    train.forceMinPn = 0
+    train.withPnBrake = False
+    train.powerLosses = get_power_loss_function(train, "static")
+    opts = {'numIntervals': 1000, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 2}, 'energyOptimal': True}
 
     SBB_track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
-    SBB_track.updateTrainLengthDependentValues(train)
+    # SBB_track.updateTrainLengthDependentValues(train)
     solver = casadiSolver(train, SBB_track, opts)
     dfSBB, statsSBB = solver.solve(duration)
 

From 654534a7bde9d1848daa878ad7e74babc4fbf286 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Tue, 2 Jun 2026 09:31:39 +0200
Subject: [PATCH 22/31] wip 3

---
 tracks/swisstopo/analyzeTracks.py | 2 +-
 trains/CH_Stadler_FLIRT_TPF.json  | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tracks/swisstopo/analyzeTracks.py b/tracks/swisstopo/analyzeTracks.py
index 7ce81fe..70a8cf1 100644
--- a/tracks/swisstopo/analyzeTracks.py
+++ b/tracks/swisstopo/analyzeTracks.py
@@ -83,7 +83,7 @@
     opts = {'numIntervals': 1000, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 2}, 'energyOptimal': True}
 
     SBB_track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
-    # SBB_track.updateTrainLengthDependentValues(train)
+    SBB_track.updateTrainLengthDependentValues(train)
     solver = casadiSolver(train, SBB_track, opts)
     dfSBB, statsSBB = solver.solve(duration)
 
diff --git a/trains/CH_Stadler_FLIRT_TPF.json b/trains/CH_Stadler_FLIRT_TPF.json
index a97c321..a0e0ff9 100644
--- a/trains/CH_Stadler_FLIRT_TPF.json
+++ b/trains/CH_Stadler_FLIRT_TPF.json
@@ -71,11 +71,11 @@
     },
     "efficiency traction": {
         "unit": "%",
-        "value": 0.9
+        "value": 90.0
     },
     "efficiency reg brake": {
         "unit": "%",
-        "value": 0.9
+        "value": 90.0
     },
     "tunnel resistance": {
         "units": {

From 297f6a9587249f7e15e6f116163d0c1dd435d3f3 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Tue, 2 Jun 2026 11:06:52 +0200
Subject: [PATCH 23/31] wip 4

---
 train.py                                      |  32 ++--
 .../gradient.py                               | 154 ++++++++++++++++++
 2 files changed, 175 insertions(+), 11 deletions(-)

diff --git a/train.py b/train.py
index bb62eae..86b945a 100644
--- a/train.py
+++ b/train.py
@@ -135,6 +135,10 @@ def __init__(self, config, pathJSON='trains') -> None:
 
     def checkFields(self):
 
+        if self.length is None or self.length < 0 or np.isinf(self.length):
+
+            raise ValueError("Train length must be a positive number, not {}!".format(self.length))
+
         if self.mass is None or self.mass < 0 or np.isinf(self.mass):
 
             raise ValueError("Train mass must be a positive number, not {}!".format(self.mass))
@@ -192,6 +196,17 @@ def checkFields(self):
                 raise ValueError("Rolling resistance coefficient {} must be positive, not {}!".format('r'+ii, coef))
 
 
+        for crossSection, coef in self.tunnelCoefficients.items():
+
+            if crossSection is None or crossSection <= 0 or np.isinf(crossSection):
+
+                raise ValueError("Tunnel cross section must be positive, not {}!".format(crossSection))
+
+            if coef is None or coef <= 0 or np.isinf(coef):
+
+                raise ValueError("Tunnel resistance coefficient must be positive, not {}!".format(coef))
+
+
     def exportModel(self):
         "Export train model (ODE and relevant train data)."
 
@@ -364,26 +379,21 @@ def __init__(self, model, solver, optsDict={}) -> None:
 
             evalPoints = [0] + [i/ns for i in range(1, ns+1)]
 
-            z0 = ca.vertcat(model.states[1], model.states[2])
+            b0 = model.states[1]
             p0 = ca.vertcat(model.controls, model.parameters)
-            zf = self.eval(z0, p0, ca.hcat(evalPoints))
+            bf = self.eval(b0, p0, ca.hcat(evalPoints))
 
             tApprox = model.states[0]
 
-            epsVelSq = 0.0001
             for idx in range(ns):
-                # vCurr = ca.sqrt(ca.fmax(zf[0, idx], epsVelSq))
-                # vNext = ca.sqrt(ca.fmax(zf[0, idx + 1], epsVelSq))
-
-                vCurr = ca.sqrt(zf[0, idx])
-                vNext = ca.sqrt(zf[0, idx + 1])
 
+                vCurr = ca.sqrt(bf[idx])
+                vNext = ca.sqrt(bf[idx+1])
                 tApprox += 2*model.parameters[-1]*(evalPoints[idx+1]-evalPoints[idx])/(vCurr + vNext)
 
-            eval = ca.vertcat(tApprox, zf[0, ns], zf[1, ns])
-
-            self.eval = ca.Function('xNxt', [model.states, ca.vertcat(model.controls, model.parameters), ca.MX.sym('dummy')], [eval])
+            eval = ca.vertcat(tApprox, bf[-1])
 
+            self.eval = ca.Function('xNxt', [model.states, ca.vertcat(model.controls, model.parameters), ca.MX.sym('ds')], [eval])
 
     def solve(self, time, velocitySquared, ds, position=0, traction=0, pnBrake=0, gradient=0, gradientLinearTerm=0, curvature=0, curvatureLinearTerm=0, tunnelFactor=0):
 
diff --git a/unitTests/trainLengthDependentTrackAttributes/gradient.py b/unitTests/trainLengthDependentTrackAttributes/gradient.py
index e5a9d4a..a9dd949 100644
--- a/unitTests/trainLengthDependentTrackAttributes/gradient.py
+++ b/unitTests/trainLengthDependentTrackAttributes/gradient.py
@@ -145,6 +145,160 @@ def test_integrator_CVODES(self):
             plt.show()
 
 
+    def test_integrator_RK(self):
+        '''
+        Track with a linearly increasing gradient over 1000 m.
+
+        The result obtained using the piecewise linear gradient model is compared
+        against a piecewise constant midpoint approximation with increasing numbers
+        of intervals.
+
+        The piecewise constant approximation should converge to the piecewise linear
+        result for both duration and final velocity.
+        '''
+
+        train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
+
+        optsDict = {'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':0, 'numSteps': 1}}
+        opts = OptionsCasadiSolver(optsDict)
+
+        trainModel = train.exportModel()
+        trainIntegrator = TrainIntegrator(trainModel, opts.integrationMethod, opts.integrationOptions.toDict())
+
+        # Scenario
+        time0 = 0
+        velSq0 = 1
+        ds = 1000
+        traction = 0.8 * (train.forceMax / train.mass)
+
+        initialGradient = 0
+        finalGradient = 0.07
+
+        maxIntervals = 50
+        relativeTolerance = 1e-3
+        plotDebug = True
+
+        # PWL gradient reference
+        pwlTimes = []
+        pwlVelocities = []
+        pwlIntervalCounts = []
+
+        for numStep in range(1, maxIntervals + 1):
+
+            optsDict = {'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1, 'numSteps': numStep}}
+            opts = OptionsCasadiSolver(optsDict)
+            pwlTrainIntegrator = TrainIntegrator(trainModel, opts.integrationMethod, opts.integrationOptions.toDict())
+
+            out = pwlTrainIntegrator.solve(
+                time=time0,
+                velocitySquared=velSq0,
+                ds=ds,
+                traction=traction,
+                pnBrake=0,
+                gradient=initialGradient,
+                gradientLinearTerm=(finalGradient - initialGradient) / ds,
+                curvature=0,
+                curvatureLinearTerm=0,
+                tunnelFactor=0
+            )
+
+            pwlTimes.append(float(out['time']))
+            pwlVelocities.append(np.sqrt(float(out['velSquared'])))
+            pwlIntervalCounts.append(numStep)
+
+        finalPwlDuration = pwlTimes[-1]
+        finalPwlVelocity = pwlVelocities[-1]
+
+        # PWC gradients using midpoint rule
+        times = []
+        velocities = []
+        intervalCounts = []
+
+        for numIntervals in range(1, maxIntervals + 1):
+
+            time = time0
+            velSq = velSq0
+
+            for idx in range(numIntervals):
+
+                gradient = (initialGradient + (idx + 0.5) * (finalGradient - initialGradient) / numIntervals)
+
+                out = trainIntegrator.solve(
+                    time=time,
+                    velocitySquared=velSq,
+                    ds=ds / numIntervals,
+                    traction=traction,
+                    pnBrake=0,
+                    gradient=gradient,
+                    gradientLinearTerm=0,
+                    curvature=0,
+                    curvatureLinearTerm=0,
+                    tunnelFactor=0
+                )
+
+                time = out['time']
+                velSq = out['velSquared']
+
+            intervalCounts.append(numIntervals)
+            times.append(float(time))
+            velocities.append(np.sqrt(float(velSq)))
+
+        finalPwcDuration = times[-1]
+        finalPwcVelocity = velocities[-1]
+
+        relativeDurationError = abs(finalPwcDuration - finalPwlDuration) / finalPwlDuration
+        relativeVelocityError = abs(finalPwcVelocity - finalPwlVelocity) / finalPwlVelocity
+
+        self.assertLess(
+            relativeDurationError,
+            relativeTolerance,
+            msg=(
+                "PWC midpoint approximation did not converge sufficiently "
+                "to the PWL duration. "
+                f"PWL duration: {finalPwlDuration:.6f}, "
+                f"PWC duration: {finalPwcDuration:.6f}, "
+                f"relative error: {relativeDurationError:.6e}."
+            )
+        )
+
+        self.assertLess(
+            relativeVelocityError,
+            relativeTolerance,
+            msg=(
+                "PWC midpoint approximation did not converge sufficiently "
+                "to the PWL final velocity. "
+                f"PWL velocity: {finalPwlVelocity:.6f}, "
+                f"PWC velocity: {finalPwcVelocity:.6f}, "
+                f"relative error: {relativeVelocityError:.6e}."
+            )
+        )
+
+        if plotDebug:
+            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(16, 10))
+
+            ax1.plot(pwlIntervalCounts, pwlTimes, marker="o", label="pwl")
+            ax1.plot(intervalCounts, times, marker="o", color="orange", label="pwc midpoint")
+            ax1.set_xlabel("Number of intervals")
+            ax1.set_ylabel("Duration [s]")
+            ax1.grid(True, which="both", linestyle="--", alpha=0.5)
+            ax1.legend(loc="upper right")
+
+            ax2.plot(pwlIntervalCounts, pwlVelocities, marker="o", label="pwl")
+            ax2.plot(intervalCounts, velocities, marker="o", color="orange", label="pwc midpoint")
+            ax2.set_xlabel("Number of intervals")
+            ax2.set_ylabel("Velocity [m/s]")
+            ax2.grid(True, which="both", linestyle="--", alpha=0.5)
+            ax2.legend(loc="upper right")
+
+            fig.tight_layout()
+            plt.show()
+
+
+    def testPWLProfile(self):
+        '''
+        Should result in same target altitude
+        '''
+
 
     def testLinearGradient(self):
         '''

From 4c8cd778a19128d83e9d22636df52101e3ac0a2e Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Tue, 2 Jun 2026 15:09:13 +0200
Subject: [PATCH 24/31] gradient unit tests finished

---
 simulations/sim_launcher.py                   |   2 +-
 track.py                                      |  16 +-
 train.py                                      |  13 +-
 .../gradient.py                               | 355 ++++++++++++++++--
 4 files changed, 337 insertions(+), 49 deletions(-)

diff --git a/simulations/sim_launcher.py b/simulations/sim_launcher.py
index 9bf3443..c08acd5 100644
--- a/simulations/sim_launcher.py
+++ b/simulations/sim_launcher.py
@@ -32,8 +32,8 @@ def get_power_loss_function(train, mode="perfect",* ,auxiliaries: float = 27_000
 
     track = Track(config={'id':'CH_ZH_LU'}, pathJSON='../tracks')
     track = Track(config={'id':'CH_StGallen_Wil'}, pathJSON='../tracks')
-    track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
     track.updateTrainLengthDependentValues(train)
+    track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
 
     opts = {'numIntervals':600, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':1}, 'energyOptimal':True}
 
diff --git a/track.py b/track.py
index 4fcfed1..29ceccd 100644
--- a/track.py
+++ b/track.py
@@ -641,6 +641,10 @@ def updateGradientsToTrainLength(self, trainLength):
 
             return
 
+        # assume train starts on a flat track
+        gradientValues = np.r_[0, gradientValues]
+        gradientPositions = np.r_[-trainLength, gradientPositions]
+
         # Each original gradient jump is spread linearly over one train length assuming uniform mass.
         # The first point has no previous gradient, so its slope contribution is zero.
         gradientJumpSlopes = np.r_[0.0, (gradientValues[1:] - gradientValues[:-1]) / trainLength]
@@ -649,8 +653,8 @@ def updateGradientsToTrainLength(self, trainLength):
         adjustedPositions = np.sort(np.unique(np.r_[gradientPositions, gradientPositions + trainLength]))
         adjustedPositions = adjustedPositions[adjustedPositions < self.length]
 
-        adjustedGradients = [gradientValues[0]]
-        gradientLinearTerms = [0.0]
+        adjustedGradients = [0]
+        gradientLinearTerms = [gradientValues[0]/trainLength]
 
         epsilon = 1e-3
 
@@ -675,7 +679,11 @@ def updateGradientsToTrainLength(self, trainLength):
             adjustedGradients.append(currentGradient)
             gradientLinearTerms.append(currentLinearTerm)
 
-        # plotGradients(self, np.asarray(pos_adj, dtype=float), np.asarray(g_adj, dtype=float), np.asarray(g_linear, dtype=float))
+        adjustedPositions = adjustedPositions[1:]
+        adjustedGradients = adjustedGradients[1:]
+        gradientLinearTerms = gradientLinearTerms[1:]
+
+        # plotGradients(self, np.asarray(adjustedPositions, dtype=float), np.asarray(adjustedGradients, dtype=float), np.asarray(gradientLinearTerms, dtype=float))
 
         self.gradients = pd.DataFrame({"Gradient [permil]": adjustedGradients, "Gradient linear term [permil/m]": gradientLinearTerms}, index=adjustedPositions)
 
@@ -695,6 +703,8 @@ def updateCurvaturesToTrainLength(self, trainLength):
 
             return
 
+         # todo: same as gradient
+
         # Each original curvature jump is spread linearly over one train length assuming uniform mass.
         # The first point has no previous curvature, so its slope contribution is zero.
         curvatureJumpSlopes = np.r_[0.0, (curvatureValues[1:] - curvatureValues[:-1]) / trainLength]
diff --git a/train.py b/train.py
index 86b945a..f419b6a 100644
--- a/train.py
+++ b/train.py
@@ -366,8 +366,6 @@ def __init__(self, model, solver, optsDict={}) -> None:
             opts = OptionsCVODES(optsDict)
             opts.numApproxSteps = 0
 
-            states = model.states
-
             t0, tf = 0, 1
             cvodesFun = ca.integrator('integrator', 'cvodes', {'x': model.states, 'p': params, 'ode': model.ode}, t0, tf, {'abstol': opts.absTol, 'reltol': opts.relTol})
 
@@ -379,22 +377,23 @@ def __init__(self, model, solver, optsDict={}) -> None:
 
             evalPoints = [0] + [i/ns for i in range(1, ns+1)]
 
-            b0 = model.states[1]
+            z0 = ca.vertcat(model.states[1], model.states[2])
             p0 = ca.vertcat(model.controls, model.parameters)
-            bf = self.eval(b0, p0, ca.hcat(evalPoints))
+            zf = self.eval(z0, p0, ca.hcat(evalPoints))  # zf[0, idx]: velSq, zf[1, idx]: pos
 
             tApprox = model.states[0]
 
             for idx in range(ns):
 
-                vCurr = ca.sqrt(bf[idx])
-                vNext = ca.sqrt(bf[idx+1])
+                vCurr = ca.sqrt(zf[0, idx])
+                vNext = ca.sqrt(zf[0, idx+1])
                 tApprox += 2*model.parameters[-1]*(evalPoints[idx+1]-evalPoints[idx])/(vCurr + vNext)
 
-            eval = ca.vertcat(tApprox, bf[-1])
+            eval = ca.vertcat(tApprox, zf[0, ns], zf[1, ns])
 
             self.eval = ca.Function('xNxt', [model.states, ca.vertcat(model.controls, model.parameters), ca.MX.sym('ds')], [eval])
 
+
     def solve(self, time, velocitySquared, ds, position=0, traction=0, pnBrake=0, gradient=0, gradientLinearTerm=0, curvature=0, curvatureLinearTerm=0, tunnelFactor=0):
 
         withPnBrake = self.model.withPnBrake
diff --git a/unitTests/trainLengthDependentTrackAttributes/gradient.py b/unitTests/trainLengthDependentTrackAttributes/gradient.py
index a9dd949..4bd7c78 100644
--- a/unitTests/trainLengthDependentTrackAttributes/gradient.py
+++ b/unitTests/trainLengthDependentTrackAttributes/gradient.py
@@ -1,4 +1,5 @@
 import unittest
+from time import perf_counter_ns
 
 import numpy as np
 from matplotlib import pyplot as plt
@@ -15,11 +16,12 @@ def test_integrator_CVODES(self):
         Track with a linearly increasing gradient over 1000 m.
 
         The result obtained using the piecewise linear gradient model is compared
-        against a piecewise constant midpoint approximation with increasing numbers
-        of intervals.
+        against a piecewise constant midpoint approximation of the gradient with increasing numbers of intervals.
 
         The piecewise constant approximation should converge to the piecewise linear
         result for both duration and final velocity.
+
+        CVODES is used as the integrator.
         '''
 
         train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
@@ -127,16 +129,18 @@ def test_integrator_CVODES(self):
         if plotDebug:
             fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(16, 10))
 
+            fig.suptitle("CVODES: PWC midpoint gradient approximation compared to PWL gradient", fontsize=14)
+
             ax1.axhline(pwlDuration, label="pwl")
             ax1.plot(intervalCounts, times, marker="o", color="orange", label="pwc midpoint")
-            ax1.set_xlabel("Number of intervals")
+            ax1.set_xlabel("Number of intervals of pwc gradient approximation")
             ax1.set_ylabel("Duration [s]")
             ax1.grid(True, which="both", linestyle="--", alpha=0.5)
             ax1.legend(loc="upper right")
 
             ax2.axhline(pwlVelocity, label="pwl")
             ax2.plot(intervalCounts, velocities, marker="o", color="orange", label="pwc midpoint")
-            ax2.set_xlabel("Number of intervals")
+            ax2.set_xlabel("Number of intervals of pwc gradient approximation")
             ax2.set_ylabel("Velocity [m/s]")
             ax2.grid(True, which="both", linestyle="--", alpha=0.5)
             ax2.legend(loc="upper right")
@@ -150,16 +154,18 @@ def test_integrator_RK(self):
         Track with a linearly increasing gradient over 1000 m.
 
         The result obtained using the piecewise linear gradient model is compared
-        against a piecewise constant midpoint approximation with increasing numbers
-        of intervals.
+        against a piecewise constant midpoint approximation of the gradient with increasing numbers of intervals.
 
         The piecewise constant approximation should converge to the piecewise linear
         result for both duration and final velocity.
+
+        RK without time approximation is used as the integrator.
+        RK substeps are increased until convergence
         '''
 
         train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
 
-        optsDict = {'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':0, 'numSteps': 1}}
+        optsDict = {'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps': 0, 'numSteps': 1}}
         opts = OptionsCasadiSolver(optsDict)
 
         trainModel = train.exportModel()
@@ -185,7 +191,7 @@ def test_integrator_RK(self):
 
         for numStep in range(1, maxIntervals + 1):
 
-            optsDict = {'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1, 'numSteps': numStep}}
+            optsDict = {'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 0, 'numSteps': numStep}}
             opts = OptionsCasadiSolver(optsDict)
             pwlTrainIntegrator = TrainIntegrator(trainModel, opts.integrationMethod, opts.integrationOptions.toDict())
 
@@ -210,9 +216,9 @@ def test_integrator_RK(self):
         finalPwlVelocity = pwlVelocities[-1]
 
         # PWC gradients using midpoint rule
-        times = []
-        velocities = []
-        intervalCounts = []
+        pwcTimes = []
+        pwcVelocities = []
+        pwcIntervalCounts = []
 
         for numIntervals in range(1, maxIntervals + 1):
 
@@ -239,12 +245,12 @@ def test_integrator_RK(self):
                 time = out['time']
                 velSq = out['velSquared']
 
-            intervalCounts.append(numIntervals)
-            times.append(float(time))
-            velocities.append(np.sqrt(float(velSq)))
+            pwcIntervalCounts.append(numIntervals)
+            pwcTimes.append(float(time))
+            pwcVelocities.append(np.sqrt(float(velSq)))
 
-        finalPwcDuration = times[-1]
-        finalPwcVelocity = velocities[-1]
+        finalPwcDuration = pwcTimes[-1]
+        finalPwcVelocity = pwcVelocities[-1]
 
         relativeDurationError = abs(finalPwcDuration - finalPwlDuration) / finalPwlDuration
         relativeVelocityError = abs(finalPwcVelocity - finalPwlVelocity) / finalPwlVelocity
@@ -276,19 +282,174 @@ def test_integrator_RK(self):
         if plotDebug:
             fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(16, 10))
 
-            ax1.plot(pwlIntervalCounts, pwlTimes, marker="o", label="pwl")
-            ax1.plot(intervalCounts, times, marker="o", color="orange", label="pwc midpoint")
-            ax1.set_xlabel("Number of intervals")
+            fig.suptitle("Explicit RK: PWL gradient with RK substeps vs PWC midpoint approximation", fontsize=14)
+
+            ax1.plot(pwlIntervalCounts, pwlTimes, marker="o", label="PWL: RK substeps")
+            ax1.plot(pwcIntervalCounts, pwcTimes, marker="o", color="orange", label="PWC: intervals")
+            ax1.set_xlabel("Refinement level: PWC intervals and RK substeps")
             ax1.set_ylabel("Duration [s]")
             ax1.grid(True, which="both", linestyle="--", alpha=0.5)
             ax1.legend(loc="upper right")
 
-            ax2.plot(pwlIntervalCounts, pwlVelocities, marker="o", label="pwl")
-            ax2.plot(intervalCounts, velocities, marker="o", color="orange", label="pwc midpoint")
-            ax2.set_xlabel("Number of intervals")
+            ax2.plot(pwlIntervalCounts, pwlVelocities, marker="o", label="PWL: RK substeps")
+            ax2.plot(pwcIntervalCounts, pwcVelocities, marker="o", color="orange", label="PWC: intervals")
+            ax2.set_xlabel("Refinement level: PWC intervals and RK substeps")
+            ax2.set_ylabel("Velocity [m/s]")
+            ax2.grid(True, which="both", linestyle="--", alpha=0.5)
+            ax2.legend(loc="upper right")
+
+            fig.tight_layout()
+            plt.show()
+
+
+    def test_integrator_RK_with_Time_Approx(self):
+        '''
+        Track with a linearly increasing gradient over 1000 m.
+
+        The result obtained using the piecewise linear gradient model is compared
+        against a piecewise constant midpoint approximation of the gradient.
+
+        The piecewise constant approximation should converge to the piecewise linear
+        result for both duration and final velocity.
+
+        RK with time approximation is used as the integrator.
+        RK uses 50 substeps.
+        Time approx steps are increased until convergence
+        '''
+
+        train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
+        trainModel = train.exportModel()
+
+        # Scenario
+        time0 = 0
+        velSq0 = 1
+        ds = 1000
+        traction = 0.8 * (train.forceMax / train.mass)
+
+        initialGradient = 0
+        finalGradient = 0.07
+
+        numIntervals = 50
+        timeApproxSteps = 30
+        relativeTolerance = 1e-3
+        plotDebug = True
+
+        # PWL gradient reference
+        pwlTimes = []
+        pwlVelocities = []
+        timeApproxStepCounts  = []
+
+        for timeSteps in range(1, timeApproxSteps + 1):
+
+            optsDict = {'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': timeSteps, 'numSteps': 50}}
+            opts = OptionsCasadiSolver(optsDict)
+            trainIntegrator = TrainIntegrator(trainModel, opts.integrationMethod, opts.integrationOptions.toDict())
+
+            out = trainIntegrator.solve(
+                time=time0,
+                velocitySquared=velSq0,
+                ds=ds,
+                traction=traction,
+                pnBrake=0,
+                gradient=initialGradient,
+                gradientLinearTerm=(finalGradient - initialGradient) / ds,
+                curvature=0,
+                curvatureLinearTerm=0,
+                tunnelFactor=0
+            )
+
+            pwlTimes.append(float(out['time']))
+            pwlVelocities.append(np.sqrt(float(out['velSquared'])))
+            timeApproxStepCounts.append(timeSteps)
+
+        finalPwlDuration = pwlTimes[-1]
+        finalPwlVelocity = pwlVelocities[-1]
+
+        # PWC gradients using midpoint rule
+        pwcTimes = []
+        pwcVelocities = []
+
+        for timeSteps in range(1, timeApproxSteps + 1):
+
+            optsDict = {'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': timeSteps, 'numSteps': 50}}
+            opts = OptionsCasadiSolver(optsDict)
+            trainIntegrator = TrainIntegrator(trainModel, opts.integrationMethod, opts.integrationOptions.toDict())
+
+            time = time0
+            velSq = velSq0
+
+            for idx in range(numIntervals):
+
+                gradient = (initialGradient + (idx + 0.5) * (finalGradient - initialGradient) / numIntervals)
+
+                out = trainIntegrator.solve(
+                    time=time,
+                    velocitySquared=velSq,
+                    ds=ds / numIntervals,
+                    traction=traction,
+                    pnBrake=0,
+                    gradient=gradient,
+                    gradientLinearTerm=0,
+                    curvature=0,
+                    curvatureLinearTerm=0,
+                    tunnelFactor=0
+                )
+
+                time = out['time']
+                velSq = out['velSquared']
+
+            pwcTimes.append(float(time))
+            pwcVelocities.append(np.sqrt(float(velSq)))
+
+        finalPwcDuration = pwcTimes[-1]
+        finalPwcVelocity = pwcVelocities[-1]
+
+        relativeDurationError = abs(finalPwcDuration - finalPwlDuration) / finalPwlDuration
+        relativeVelocityError = abs(finalPwcVelocity - finalPwlVelocity) / finalPwlVelocity
+
+        self.assertLess(
+            relativeDurationError,
+            relativeTolerance,
+            msg=(
+                "PWC midpoint approximation did not converge sufficiently "
+                "to the PWL duration. "
+                f"PWL duration: {finalPwlDuration:.6f}, "
+                f"PWC duration: {finalPwcDuration:.6f}, "
+                f"relative error: {relativeDurationError:.6e}."
+            )
+        )
+
+        self.assertLess(
+            relativeVelocityError,
+            relativeTolerance,
+            msg=(
+                "PWC midpoint approximation did not converge sufficiently "
+                "to the PWL final velocity. "
+                f"PWL velocity: {finalPwlVelocity:.6f}, "
+                f"PWC velocity: {finalPwcVelocity:.6f}, "
+                f"relative error: {relativeVelocityError:.6e}."
+            )
+        )
+
+        if plotDebug:
+            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(16, 10))
+
+            fig.suptitle("RK with Time Approx: PWC midpoint gradient approximation compared to PWL gradient", fontsize=14)
+
+            ax1.plot(timeApproxStepCounts , pwlTimes, marker="o", label="pwl")
+            ax1.plot(timeApproxStepCounts, pwcTimes, marker="o", color="orange", label="pwc midpoint")
+            ax1.set_xlabel("Number of time approx steps")
+            ax1.set_ylabel("Duration [s]")
+            ax1.grid(True, which="both", linestyle="--", alpha=0.5)
+            ax1.legend(loc="upper right")
+
+            ax2.plot(timeApproxStepCounts , pwlVelocities, marker="o", label="pwl")
+            ax2.plot(timeApproxStepCounts, pwcVelocities, marker="o", color="orange", label="pwc midpoint")
+            ax2.set_xlabel("Number of time approx steps")
             ax2.set_ylabel("Velocity [m/s]")
             ax2.grid(True, which="both", linestyle="--", alpha=0.5)
             ax2.legend(loc="upper right")
+            ax2.ticklabel_format(axis="y", style="plain", useOffset=False)
 
             fig.tight_layout()
             plt.show()
@@ -296,9 +457,77 @@ def test_integrator_RK(self):
 
     def testPWLProfile(self):
         '''
-        Should result in same target altitude
+        Compare the final altitude of the original length-independent gradient profile
+        with the train-length-dependent piecewise linear profile.
+
+        Both profiles should start from the same altitude and end at the same target altitude.
         '''
 
+        plotDebug = True
+        altitudeTolerance = 1e-6
+
+        trainLength = 800   # [m]
+        track = Track(config={'id': 'CH_StGallen_Wil'}, pathJSON='tracks')
+
+        # track needs to be flat at least train length meters before the end of the track
+        track.gradients = track.gradients[track.gradients.index < track.length - trainLength]
+        track.gradients.loc[track.length - trainLength] = {"Gradient [permil]": 0.0, "Gradient linear term [permil/m]": 0.0}
+
+        df_alt = computeAltitude(track.gradients, track.length)
+
+        train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
+        train.length = trainLength
+        track.updateTrainLengthDependentValues(train)
+
+        positions = track.gradients.index.to_numpy(dtype=float)
+        positions = np.r_[positions, track.length]
+
+        gradient = track.gradients["Gradient [permil]"].to_numpy(dtype=float)
+        gradientLinear = track.gradients["Gradient linear term [permil/m]"].to_numpy(dtype=float)
+
+        ds = positions[1:] - positions[:-1]
+
+        pwlAltitude = np.zeros(len(positions))
+        pwlAltitude[1:] = np.cumsum((gradient * ds + 0.5 * gradientLinear * ds ** 2) / 1000)
+
+        self.assertAlmostEqual(
+            df_alt["Altitude [m]"].iloc[0],
+            pwlAltitude[0],
+            delta=altitudeTolerance,
+            msg="Length-independent and train-length-dependent altitude profiles should start at the same altitude."
+        )
+
+        lengthIndependentFinalAltitude = df_alt["Altitude [m]"].iloc[-1]
+        lengthDependentFinalAltitude = pwlAltitude[-1]
+
+        self.assertAlmostEqual(
+            lengthIndependentFinalAltitude,
+            lengthDependentFinalAltitude,
+            delta=altitudeTolerance,
+            msg=(
+                "Length-independent and train-length-dependent altitude profiles "
+                "should end at the same altitude. "
+                f"Length-independent final altitude: {lengthIndependentFinalAltitude:.8f} m, "
+                f"train-length-dependent final altitude: {lengthDependentFinalAltitude:.8f} m."
+            )
+        )
+
+        if plotDebug:
+
+            fig, ax = plt.subplots(figsize=(16, 8))
+
+            ax.plot(df_alt.index.values / 1000, df_alt["Altitude [m]"].to_numpy(), label="length-independent altitude")
+            ax.plot(positions / 1000, pwlAltitude, label="length-dependent altitude")
+            ax.set_title("Altitude Comparison")
+            ax.set_xlabel("Position [km]")
+            ax.set_ylabel("Altitude [m]")
+            ax.grid(True, which="both", linestyle="--", alpha=0.5)
+            ax.legend(loc="upper right")
+            ax.set_xlim(0, track.length / 1000)
+            ax.figure.tight_layout()
+
+            plt.show()
+
 
     def testLinearGradient(self):
         '''
@@ -306,7 +535,7 @@ def testLinearGradient(self):
         decrease from 3000 m to 4000 m.
 
         Energy consumption should be roughly equal if computed using piecewise
-        linear gradients or equivalent piecewise constant gradients.
+        linear gradients or equivalent piecewise constant gradients due to a high number of shooting intervals.
 
         Altitude should be 0 m at the end.
         '''
@@ -316,8 +545,8 @@ def testLinearGradient(self):
         duration = 5000/(60/3.6) # [s]
 
         altitudeTolerance = 1e-4
-        energyRelativeTolerance = 1e-4
-        numIntervals = 100
+        energyRelativeTolerance = 1e-3
+        numIntervals = 50
 
         train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
         train.length = 600
@@ -326,23 +555,25 @@ def testLinearGradient(self):
         track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
         track.updateTrainLengthDependentValues(train)
 
-        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 0}, 'energyOptimal': True}
+        # PWL Gradients
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'CVODES'}
         solver = casadiSolver(train, track, opts)
-        df1, stats1 = solver.solve(duration)
+        pwl_df, pwl_stats = solver.solve(duration)
 
-        energyConsumptionWithLinearTerms = stats1['Cost']
+        energyConsumptionWithLinearTerms = pwl_stats['Cost']
 
-        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 0}, 'energyOptimal': True, 'pwcLengthDependentTrackAttributes': True}
+        # PWC Gradients
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'CVODES', 'pwcLengthDependentTrackAttributes': True}
         solver = casadiSolver(train, track, opts)
-        df2, stats2 = solver.solve(duration)
+        pwc_df, pwc_stats = solver.solve(duration)
 
-        energyConsumptionWithPwcTerms= stats2['Cost']
+        energyConsumptionWithPwcTerms= pwc_stats['Cost']
 
         relativeEnergyDifference = (abs(energyConsumptionWithLinearTerms - energyConsumptionWithPwcTerms) / energyConsumptionWithLinearTerms)
 
-        df_grads_2 = df2.set_index("Position [m]")[["Gradient [permil]"]]
-        df_alt_2 = computeAltitude(df_grads_2, track.length)
-        finalAltitude = df_alt_2.iloc[-1]["Altitude [m]"]
+        pwc_df_grads = pwc_df.set_index("Position [m]")[["Gradient [permil]"]]
+        pwc_df_alt = computeAltitude(pwc_df_grads, track.length)
+        finalAltitude = pwc_df_alt.iloc[-1]["Altitude [m]"]
 
         self.assertLess(
             relativeEnergyDifference,
@@ -372,9 +603,9 @@ def testLinearGradient(self):
 
             fig, ax = plt.subplots(figsize=(16, 8))
 
-            df_grads_1 = df1.set_index("Position [m]")[["Gradient [permil]"]]
+            df_grads_1 = pwl_df.set_index("Position [m]")[["Gradient [permil]"]]
             ax.plot(df_grads_1.index.values / 1000, df_grads_1["Gradient [permil]"],label="pwl gradients")
-            ax.step(df_grads_2.index.values / 1000, df_grads_2["Gradient [permil]"], '--', where='post', label="pwc gradients")
+            ax.step(pwc_df_grads.index.values / 1000, pwc_df_grads["Gradient [permil]"], '--', where='post', label="pwc gradients")
             ax.set_title("Gradients")
             ax.set_xlabel("Position [km]")
             ax.set_ylabel("Gradient [‰]")
@@ -383,4 +614,52 @@ def testLinearGradient(self):
             ax.set_xlim(0, track.length / 1000)
             ax.figure.tight_layout()
 
-            plt.show()
\ No newline at end of file
+            plt.show()
+
+
+    def testSmallTrainLengthNotAffectingEnergyConsumption(self):
+        '''
+        Use an artificially short train on a real track profile.
+
+        For a very small train length, the train-length-dependent gradient profile
+        should be almost identical to the original gradient profile. Therefore, the
+        energy consumption should remain within a small relative tolerance.
+        '''
+
+        relativeTolerance = 0.01
+        trainLength = 10  # [m]
+        train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
+        train.length = trainLength
+
+        track = Track(config={'id': 'CH_StGallen_Wil'}, pathJSON='tracks')
+
+        # track needs to be flat at least train length meters before the end of the track
+        track.gradients = track.gradients[track.gradients.index < track.length - trainLength]
+        track.gradients.loc[track.length - trainLength] = {"Gradient [permil]": 0.0, "Gradient linear term [permil/m]": 0.0}
+
+        duration = track.length / (80/3.6)
+
+        # train-length-independent
+        opts = {'numIntervals':600, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':0}, 'energyOptimal':True}
+        solver = casadiSolver(train, track, opts)
+        indep_df, indep_stats = solver.solve(duration)
+
+        track.updateTrainLengthDependentValues(train)
+        solver = casadiSolver(train, track, opts)
+        dep_df, dep_stats = solver.solve(duration)
+
+        energyConsumptionIndependentOfTrainLength = indep_stats['Cost']
+        energyConsumptionDependentOfTrainLength = dep_stats['Cost']
+
+        relativeDifference = (abs(energyConsumptionDependentOfTrainLength - energyConsumptionIndependentOfTrainLength) / energyConsumptionIndependentOfTrainLength)
+
+        self.assertLess(
+            relativeDifference,
+            relativeTolerance,
+            msg=(
+                "Energy consumption with and without train-length-dependent gradients should be roughly equal. "
+                f"Independent: {energyConsumptionIndependentOfTrainLength:.6f}, "
+                f"dependent: {energyConsumptionDependentOfTrainLength:.6f}, "
+                f"relative difference: {relativeDifference:.6e}."
+            )
+        )
\ No newline at end of file

From cc1dbcb221268814a28d77ce9afcd63de6097403 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Wed, 3 Jun 2026 11:41:00 +0200
Subject: [PATCH 25/31] unit test amended

---
 track.py                                      |   23 +-
 tracks/CH_StGallen_Wil.json                   | 1203 ++++++++++++++++-
 unitTests/integrators/integrators.py          |    8 +-
 .../curvature.py                              |  671 +++++++--
 .../gradient.py                               |   34 +-
 5 files changed, 1819 insertions(+), 120 deletions(-)

diff --git a/track.py b/track.py
index 29ceccd..f8a0c3c 100644
--- a/track.py
+++ b/track.py
@@ -8,6 +8,8 @@
 from utils import checkTTOBenchVersion, convertUnit, pickEquallySpacedPoints, plotSpeedLimits, plotGradients, \
     plotCurvatures
 
+plotDebug = False
+
 
 def importTuples(tuples, xLabel, yLabels):
     """
@@ -352,6 +354,7 @@ def importCurvatureTuples(self, tuples, unitRadiusStart='m', unitRadiusEnd='m',
         tuples = self.sampleClothoid(tuples, clothoidSamplingInterval)
 
         self.curvatures = importTuples(tuples, 'Position [m]', ['Curvature [1/m]'])
+        self.curvatures["Curvature [1/m]"] = self.curvatures["Curvature [1/m]"].abs()
 
         checkDataFrame(self.curvatures, self.length)
 
@@ -618,7 +621,9 @@ def updateSpeedLimitsToTrainLength(self, trainLength):
                 pos_adj.append(new_pos)
                 v_adj.append(v[i])
 
-            # plotSpeedLimits(self, np.asarray(pos_adj, dtype=float), np.asarray(v_adj, dtype=float))
+            if plotDebug:
+
+                plotSpeedLimits(self, np.asarray(pos_adj, dtype=float), np.asarray(v_adj, dtype=float))
 
             self.speedLimits = pd.DataFrame(
                 {"Speed limit [m/s]": v_adj},
@@ -683,7 +688,9 @@ def updateGradientsToTrainLength(self, trainLength):
         adjustedGradients = adjustedGradients[1:]
         gradientLinearTerms = gradientLinearTerms[1:]
 
-        # plotGradients(self, np.asarray(adjustedPositions, dtype=float), np.asarray(adjustedGradients, dtype=float), np.asarray(gradientLinearTerms, dtype=float))
+        if plotDebug:
+
+            plotGradients(self, np.asarray(adjustedPositions, dtype=float), np.asarray(adjustedGradients, dtype=float), np.asarray(gradientLinearTerms, dtype=float))
 
         self.gradients = pd.DataFrame({"Gradient [permil]": adjustedGradients, "Gradient linear term [permil/m]": gradientLinearTerms}, index=adjustedPositions)
 
@@ -703,7 +710,9 @@ def updateCurvaturesToTrainLength(self, trainLength):
 
             return
 
-         # todo: same as gradient
+        # assume train starts on a straight track
+        curvatureValues = np.r_[0, curvatureValues]
+        curvaturePositions = np.r_[-trainLength, curvaturePositions]
 
         # Each original curvature jump is spread linearly over one train length assuming uniform mass.
         # The first point has no previous curvature, so its slope contribution is zero.
@@ -740,7 +749,13 @@ def updateCurvaturesToTrainLength(self, trainLength):
             adjustedCurvatures.append(currentCurvature)
             curvatureLinearTerms.append(currentLinearTerm)
 
-        # plotCurvatures(self, np.asarray(pos_adj, dtype=float), np.asarray(c_adj, dtype=float), np.asarray(c_linear, dtype=float))
+        adjustedPositions = adjustedPositions[1:]
+        adjustedCurvatures = adjustedCurvatures[1:]
+        curvatureLinearTerms = curvatureLinearTerms[1:]
+
+        if plotDebug:
+
+            plotCurvatures(self, np.asarray(adjustedPositions, dtype=float), np.asarray(adjustedCurvatures, dtype=float), np.asarray(curvatureLinearTerms, dtype=float))
 
         self.curvatures = pd.DataFrame({"Curvature [1/m]": adjustedCurvatures, "Curvature linear term [1/m^2]": curvatureLinearTerms}, index=adjustedPositions)
 
diff --git a/tracks/CH_StGallen_Wil.json b/tracks/CH_StGallen_Wil.json
index 5d4199c..c2afc63 100644
--- a/tracks/CH_StGallen_Wil.json
+++ b/tracks/CH_StGallen_Wil.json
@@ -1,8 +1,8 @@
 {
     "metadata": {
         "id": "CH_StGallen_Wil",
-        "created by": "Dimitris Kouzoupis",
-        "library version": "TTOBench v1.1",
+        "created by": "Dimitris Kouzoupis, Juxhino Kavaja",
+        "library version": "TTOBench v1.2",
         "license": "BSD 2-Clause License"
     },
     "altitude": {
@@ -695,5 +695,1204 @@
                 -4.6
             ]
         ]
+    },
+    "curvatures": {
+        "units": {
+            "position": "m",
+            "radius at start": "m",
+            "radius at end": "m"
+        },
+        "values": [
+            [
+                0.0,
+                502.0,
+                502.0
+            ],
+            [
+                49.6,
+                502.0,
+                3570.0
+            ],
+            [
+                125.6,
+                3570.0,
+                3570.0
+            ],
+            [
+                172.5,
+                3570.0,
+                1250.0
+            ],
+            [
+                198.5,
+                1250.0,
+                1250.0
+            ],
+            [
+                232.1,
+                1250.0,
+                "infinity"
+            ],
+            [
+                287.1,
+                "infinity",
+                "infinity"
+            ],
+            [
+                330.2,
+                -5700.0,
+                -5700.0
+            ],
+            [
+                393.1,
+                "infinity",
+                "infinity"
+            ],
+            [
+                445.4,
+                "infinity",
+                1567.0
+            ],
+            [
+                594.4,
+                1567.0,
+                1567.0
+            ],
+            [
+                948.8,
+                1567.0,
+                "infinity"
+            ],
+            [
+                1018.8,
+                "infinity",
+                -850.0
+            ],
+            [
+                1106.1,
+                -850.0,
+                -850.0
+            ],
+            [
+                1234.7,
+                -850.0,
+                -2600.0
+            ],
+            [
+                1314.7,
+                -2600.0,
+                -2600.0
+            ],
+            [
+                1425.1,
+                -2600.0,
+                "infinity"
+            ],
+            [
+                1505.1,
+                "infinity",
+                "infinity"
+            ],
+            [
+                2140.4,
+                "infinity",
+                508.0
+            ],
+            [
+                2235.8,
+                508.0,
+                508.0
+            ],
+            [
+                2326.0,
+                508.0,
+                "infinity"
+            ],
+            [
+                2422.0,
+                "infinity",
+                "infinity"
+            ],
+            [
+                2552.6,
+                "infinity",
+                -752.0
+            ],
+            [
+                2627.6,
+                -752.0,
+                -752.0
+            ],
+            [
+                2782.3,
+                -752.0,
+                -605.0
+            ],
+            [
+                2812.3,
+                -605.0,
+                -605.0
+            ],
+            [
+                2953.0,
+                -605.0,
+                "infinity"
+            ],
+            [
+                3043.0,
+                "infinity",
+                "infinity"
+            ],
+            [
+                3103.0,
+                "infinity",
+                520.2
+            ],
+            [
+                3195.0,
+                520.2,
+                520.2
+            ],
+            [
+                3380.8,
+                520.2,
+                "infinity"
+            ],
+            [
+                3475.8,
+                "infinity",
+                "infinity"
+            ],
+            [
+                3557.0,
+                "infinity",
+                2896.2
+            ],
+            [
+                3597.0,
+                2896.2,
+                2896.2
+            ],
+            [
+                3705.7,
+                2896.2,
+                6246.2
+            ],
+            [
+                3755.7,
+                6246.2,
+                6246.2
+            ],
+            [
+                3790.6,
+                6246.2,
+                1020.0
+            ],
+            [
+                3884.6,
+                1020.0,
+                1020.0
+            ],
+            [
+                3982.6,
+                1020.0,
+                "infinity"
+            ],
+            [
+                4058.6,
+                "infinity",
+                "infinity"
+            ],
+            [
+                4499.2,
+                "infinity",
+                657.2
+            ],
+            [
+                4609.2,
+                657.2,
+                657.2
+            ],
+            [
+                4717.1,
+                657.2,
+                "infinity"
+            ],
+            [
+                4824.1,
+                "infinity",
+                "infinity"
+            ],
+            [
+                4926.7,
+                "infinity",
+                5000.0
+            ],
+            [
+                4943.7,
+                5000.0,
+                5000.0
+            ],
+            [
+                4977.4,
+                5000.0,
+                "infinity"
+            ],
+            [
+                4994.4,
+                "infinity",
+                -5000.0
+            ],
+            [
+                5011.4,
+                -5000.0,
+                -5000.0
+            ],
+            [
+                5045.0,
+                -5000.0,
+                "infinity"
+            ],
+            [
+                5062.0,
+                "infinity",
+                "infinity"
+            ],
+            [
+                5486.3,
+                "infinity",
+                -6000.0
+            ],
+            [
+                5514.3,
+                -6000.0,
+                -6000.0
+            ],
+            [
+                5537.9,
+                -6000.0,
+                "infinity"
+            ],
+            [
+                5565.9,
+                "infinity",
+                6000.0
+            ],
+            [
+                5593.9,
+                6000.0,
+                6000.0
+            ],
+            [
+                5617.5,
+                6000.0,
+                "infinity"
+            ],
+            [
+                5645.5,
+                "infinity",
+                "infinity"
+            ],
+            [
+                7385.4,
+                "infinity",
+                -3003.8
+            ],
+            [
+                7425.5,
+                -3003.8,
+                -3003.8
+            ],
+            [
+                7496.1,
+                -3003.8,
+                "infinity"
+            ],
+            [
+                7536.2,
+                "infinity",
+                "infinity"
+            ],
+            [
+                7793.9,
+                "infinity",
+                -962.0
+            ],
+            [
+                7894.7,
+                -962.0,
+                -962.0
+            ],
+            [
+                8105.6,
+                -962.0,
+                "infinity"
+            ],
+            [
+                8194.0,
+                "infinity",
+                "infinity"
+            ],
+            [
+                8464.2,
+                "infinity",
+                2800.0
+            ],
+            [
+                8507.1,
+                2800.0,
+                2800.0
+            ],
+            [
+                8582.9,
+                2800.0,
+                "infinity"
+            ],
+            [
+                8625.8,
+                "infinity",
+                -2780.0
+            ],
+            [
+                8657.8,
+                -2780.0,
+                -2780.0
+            ],
+            [
+                8737.2,
+                -2780.0,
+                "infinity"
+            ],
+            [
+                8784.2,
+                "infinity",
+                "infinity"
+            ],
+            [
+                8974.8,
+                "infinity",
+                -1350.0
+            ],
+            [
+                9034.8,
+                -1350.0,
+                -1350.0
+            ],
+            [
+                9079.3,
+                -1350.0,
+                -6130.8
+            ],
+            [
+                9135.4,
+                -6130.8,
+                "infinity"
+            ],
+            [
+                9151.3,
+                "infinity",
+                "infinity"
+            ],
+            [
+                9212.0,
+                "infinity",
+                -803.6
+            ],
+            [
+                9310.0,
+                -803.6,
+                -803.6
+            ],
+            [
+                9341.8,
+                -803.6,
+                "infinity"
+            ],
+            [
+                9439.8,
+                "infinity",
+                "infinity"
+            ],
+            [
+                9629.7,
+                "infinity",
+                -778.0
+            ],
+            [
+                9734.7,
+                -778.0,
+                -778.0
+            ],
+            [
+                9943.2,
+                -778.0,
+                "infinity"
+            ],
+            [
+                10048.2,
+                "infinity",
+                "infinity"
+            ],
+            [
+                10336.8,
+                "infinity",
+                700.0
+            ],
+            [
+                10476.8,
+                700.0,
+                700.0
+            ],
+            [
+                10986.5,
+                700.0,
+                "infinity"
+            ],
+            [
+                11126.5,
+                "infinity",
+                "infinity"
+            ],
+            [
+                11501.9,
+                "infinity",
+                -6800.0
+            ],
+            [
+                11531.9,
+                -6800.0,
+                -6800.0
+            ],
+            [
+                11618.6,
+                -6800.0,
+                "infinity"
+            ],
+            [
+                11648.6,
+                "infinity",
+                "infinity"
+            ],
+            [
+                11714.3,
+                "infinity",
+                -6600.0
+            ],
+            [
+                11744.3,
+                -6600.0,
+                -6600.0
+            ],
+            [
+                11838.2,
+                -6600.0,
+                "infinity"
+            ],
+            [
+                11868.2,
+                "infinity",
+                "infinity"
+            ],
+            [
+                12241.3,
+                "infinity",
+                736.3
+            ],
+            [
+                12350.3,
+                736.3,
+                736.3
+            ],
+            [
+                12410.9,
+                736.3,
+                "infinity"
+            ],
+            [
+                12519.9,
+                "infinity",
+                "infinity"
+            ],
+            [
+                13645.7,
+                "infinity",
+                585.0
+            ],
+            [
+                13751.7,
+                585.0,
+                585.0
+            ],
+            [
+                13863.8,
+                585.0,
+                "infinity"
+            ],
+            [
+                13969.8,
+                "infinity",
+                "infinity"
+            ],
+            [
+                14000.9,
+                "infinity",
+                -5000.0
+            ],
+            [
+                14020.9,
+                -5000.0,
+                -5000.0
+            ],
+            [
+                14111.9,
+                -5000.0,
+                "infinity"
+            ],
+            [
+                14141.9,
+                "infinity",
+                5000.0
+            ],
+            [
+                14171.9,
+                5000.0,
+                5000.0
+            ],
+            [
+                14231.1,
+                5000.0,
+                "infinity"
+            ],
+            [
+                14251.1,
+                "infinity",
+                "infinity"
+            ],
+            [
+                14351.7,
+                "infinity",
+                -3000.0
+            ],
+            [
+                14381.7,
+                -3000.0,
+                -3000.0
+            ],
+            [
+                14456.5,
+                -3000.0,
+                "infinity"
+            ],
+            [
+                14486.5,
+                "infinity",
+                3000.0
+            ],
+            [
+                14516.5,
+                3000.0,
+                3000.0
+            ],
+            [
+                14591.6,
+                3000.0,
+                "infinity"
+            ],
+            [
+                14621.6,
+                "infinity",
+                "infinity"
+            ],
+            [
+                14796.0,
+                "infinity",
+                -610.4
+            ],
+            [
+                14906.0,
+                -610.4,
+                -610.4
+            ],
+            [
+                14935.8,
+                -634.8,
+                -634.8
+            ],
+            [
+                15252.8,
+                -634.8,
+                -583.8
+            ],
+            [
+                15283.0,
+                -583.8,
+                -583.8
+            ],
+            [
+                15424.4,
+                -583.8,
+                "infinity"
+            ],
+            [
+                15506.7,
+                "infinity",
+                "infinity"
+            ],
+            [
+                15625.2,
+                "infinity",
+                488.2
+            ],
+            [
+                15724.8,
+                488.2,
+                488.2
+            ],
+            [
+                15812.8,
+                488.2,
+                498.1
+            ],
+            [
+                15837.6,
+                498.1,
+                498.1
+            ],
+            [
+                16256.9,
+                498.1,
+                "infinity"
+            ],
+            [
+                16346.6,
+                "infinity",
+                "infinity"
+            ],
+            [
+                16521.0,
+                "infinity",
+                -588.8
+            ],
+            [
+                16603.3,
+                -588.8,
+                -588.8
+            ],
+            [
+                16799.4,
+                -588.8,
+                -610.8
+            ],
+            [
+                16824.6,
+                -610.8,
+                -610.8
+            ],
+            [
+                17018.6,
+                -610.8,
+                "infinity"
+            ],
+            [
+                17100.8,
+                "infinity",
+                "infinity"
+            ],
+            [
+                17718.9,
+                "infinity",
+                -533.8
+            ],
+            [
+                17811.2,
+                -533.8,
+                -533.8
+            ],
+            [
+                17860.7,
+                -533.8,
+                "infinity"
+            ],
+            [
+                17957.8,
+                "infinity",
+                "infinity"
+            ],
+            [
+                18042.0,
+                "infinity",
+                526.2
+            ],
+            [
+                18133.6,
+                526.2,
+                526.2
+            ],
+            [
+                18200.7,
+                526.2,
+                "infinity"
+            ],
+            [
+                18298.3,
+                "infinity",
+                "infinity"
+            ],
+            [
+                18438.1,
+                "infinity",
+                -591.8
+            ],
+            [
+                18534.4,
+                -591.8,
+                -591.8
+            ],
+            [
+                18658.1,
+                -591.8,
+                "infinity"
+            ],
+            [
+                18738.3,
+                "infinity",
+                "infinity"
+            ],
+            [
+                18787.4,
+                "infinity",
+                347.8
+            ],
+            [
+                18864.4,
+                347.8,
+                347.8
+            ],
+            [
+                19201.7,
+                350.0,
+                350.0
+            ],
+            [
+                19223.1,
+                343.8,
+                343.8
+            ],
+            [
+                19355.8,
+                343.8,
+                400.0
+            ],
+            [
+                19385.8,
+                400.0,
+                400.0
+            ],
+            [
+                19477.6,
+                400.0,
+                "infinity"
+            ],
+            [
+                19549.6,
+                "infinity",
+                "infinity"
+            ],
+            [
+                19655.4,
+                "infinity",
+                -403.1
+            ],
+            [
+                19752.1,
+                -403.1,
+                -403.1
+            ],
+            [
+                19845.4,
+                -388.6,
+                -388.6
+            ],
+            [
+                19870.3,
+                -399.3,
+                -399.3
+            ],
+            [
+                19974.2,
+                -399.3,
+                -1900.0
+            ],
+            [
+                20056.0,
+                -1900.0,
+                -1900.0
+            ],
+            [
+                20093.6,
+                -1900.0,
+                -700.0
+            ],
+            [
+                20138.6,
+                -700.0,
+                -700.0
+            ],
+            [
+                20183.2,
+                -700.0,
+                -393.8
+            ],
+            [
+                20228.2,
+                -393.8,
+                -393.8
+            ],
+            [
+                20473.0,
+                -393.8,
+                "infinity"
+            ],
+            [
+                20563.4,
+                "infinity",
+                "infinity"
+            ],
+            [
+                20618.0,
+                "infinity",
+                650.0
+            ],
+            [
+                20717.8,
+                650.0,
+                650.0
+            ],
+            [
+                20915.1,
+                650.0,
+                "infinity"
+            ],
+            [
+                21014.8,
+                "infinity",
+                "infinity"
+            ],
+            [
+                21375.7,
+                "infinity",
+                864.2
+            ],
+            [
+                21460.5,
+                864.2,
+                864.2
+            ],
+            [
+                21617.7,
+                864.2,
+                "infinity"
+            ],
+            [
+                21702.6,
+                "infinity",
+                "infinity"
+            ],
+            [
+                21801.7,
+                "infinity",
+                -1103.8
+            ],
+            [
+                21881.8,
+                -1103.8,
+                -1103.8
+            ],
+            [
+                22036.8,
+                -1103.8,
+                "infinity"
+            ],
+            [
+                22116.9,
+                "infinity",
+                "infinity"
+            ],
+            [
+                22496.4,
+                "infinity",
+                -645.8
+            ],
+            [
+                22586.7,
+                -645.8,
+                -645.8
+            ],
+            [
+                22814.7,
+                -645.8,
+                "infinity"
+            ],
+            [
+                22904.9,
+                "infinity",
+                "infinity"
+            ],
+            [
+                23114.4,
+                "infinity",
+                -538.8
+            ],
+            [
+                23224.8,
+                -538.8,
+                -538.8
+            ],
+            [
+                23403.5,
+                -538.8,
+                "infinity"
+            ],
+            [
+                23513.9,
+                "infinity",
+                "infinity"
+            ],
+            [
+                23932.4,
+                "infinity",
+                576.2
+            ],
+            [
+                24029.0,
+                576.2,
+                576.2
+            ],
+            [
+                24120.7,
+                576.2,
+                "infinity"
+            ],
+            [
+                24217.4,
+                "infinity",
+                "infinity"
+            ],
+            [
+                24287.7,
+                "infinity",
+                -585.8
+            ],
+            [
+                24384.0,
+                -585.8,
+                -585.8
+            ],
+            [
+                24606.4,
+                -585.8,
+                "infinity"
+            ],
+            [
+                24702.7,
+                "infinity",
+                "infinity"
+            ],
+            [
+                24776.0,
+                "infinity",
+                8400.0
+            ],
+            [
+                24816.0,
+                8400.0,
+                8400.0
+            ],
+            [
+                24879.3,
+                8400.0,
+                "infinity"
+            ],
+            [
+                24919.3,
+                "infinity",
+                "infinity"
+            ],
+            [
+                25179.4,
+                "infinity",
+                525.0
+            ],
+            [
+                25269.4,
+                525.0,
+                525.0
+            ],
+            [
+                25703.3,
+                525.0,
+                "infinity"
+            ],
+            [
+                25793.3,
+                "infinity",
+                "infinity"
+            ],
+            [
+                26242.1,
+                "infinity",
+                -491.0
+            ],
+            [
+                26336.1,
+                -491.0,
+                -491.0
+            ],
+            [
+                26588.8,
+                -491.0,
+                "infinity"
+            ],
+            [
+                26688.8,
+                "infinity",
+                "infinity"
+            ],
+            [
+                26752.4,
+                "infinity",
+                457.2
+            ],
+            [
+                26852.4,
+                457.2,
+                457.2
+            ],
+            [
+                27208.3,
+                457.2,
+                461.0
+            ],
+            [
+                27232.3,
+                461.0,
+                461.0
+            ],
+            [
+                27546.7,
+                461.0,
+                453.0
+            ],
+            [
+                27570.7,
+                453.0,
+                453.0
+            ],
+            [
+                27731.2,
+                453.0,
+                "infinity"
+            ],
+            [
+                27827.2,
+                "infinity",
+                "infinity"
+            ],
+            [
+                28456.1,
+                "infinity",
+                -372.1
+            ],
+            [
+                28532.1,
+                -372.1,
+                -372.1
+            ],
+            [
+                28638.4,
+                -372.1,
+                -397.8
+            ],
+            [
+                28678.4,
+                -397.8,
+                -397.8
+            ],
+            [
+                28748.8,
+                -397.8,
+                -340.1
+            ],
+            [
+                28788.8,
+                -340.1,
+                -340.1
+            ],
+            [
+                28879.4,
+                -340.1,
+                -520.0
+            ],
+            [
+                28919.4,
+                -520.0,
+                -520.0
+            ],
+            [
+                28948.8,
+                -520.0,
+                -351.0
+            ],
+            [
+                28988.8,
+                -351.0,
+                -351.0
+            ],
+            [
+                29090.0,
+                -351.0,
+                "infinity"
+            ],
+            [
+                29160.0,
+                "infinity",
+                "infinity"
+            ],
+            [
+                29311.1,
+                "infinity",
+                -1600.0
+            ],
+            [
+                29331.1,
+                -1600.0,
+                -1600.0
+            ],
+            [
+                29383.3,
+                -1600.0,
+                "infinity"
+            ],
+            [
+                29403.3,
+                "infinity",
+                "infinity"
+            ],
+            [
+                29457.2,
+                "infinity",
+                -490.0
+            ],
+            [
+                29507.2,
+                -490.0,
+                -490.0
+            ],
+            [
+                29531.0,
+                -490.0,
+                -901.4
+            ]
+        ]
     }
 }
\ No newline at end of file
diff --git a/unitTests/integrators/integrators.py b/unitTests/integrators/integrators.py
index f6fa23d..649802a 100644
--- a/unitTests/integrators/integrators.py
+++ b/unitTests/integrators/integrators.py
@@ -21,13 +21,13 @@ def testAllIntegratorTypesWork(self):
         endPosition = 5000  # [m]
         duration = 5000 / (60 / 3.6)  # [s]
 
-        tol = 0.02
-        numIntervals = 100
+        tol = 0.1
+        numIntervals = 200
 
-        train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='trains')
+        train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
         train.length = 600
 
-        track = Track(config={'id': 'test_one_hill'}, pathJSON='tracks')
+        track = Track(config={'id': 'CH_StGallen_Wil'}, pathJSON='tracks')
         track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
         track.updateTrainLengthDependentValues(train)
 
diff --git a/unitTests/trainLengthDependentTrackAttributes/curvature.py b/unitTests/trainLengthDependentTrackAttributes/curvature.py
index 4feb44e..932d2b0 100644
--- a/unitTests/trainLengthDependentTrackAttributes/curvature.py
+++ b/unitTests/trainLengthDependentTrackAttributes/curvature.py
@@ -1,15 +1,544 @@
 import unittest
 
+import numpy as np
+import pandas as pd
 from matplotlib import pyplot as plt
 
-from ocp import casadiSolver
-from track import Track, computeAltitude
-from train import Train
+from ocp import casadiSolver, OptionsCasadiSolver
+from track import Track
+from train import Train, TrainIntegrator
+
+
+plotDebug = False
+
+
+def computeHeadingFromCurvature(curvatures, trackLength):
+
+    positions = curvatures.index.to_numpy(dtype=float)
+
+    if positions[-1] < trackLength:
+        positions = np.r_[positions, trackLength]
+
+    curvature = curvatures["Curvature [1/m]"].to_numpy(dtype=float)
+
+    if "Curvature linear term [1/m^2]" in curvatures.columns:
+
+        curvatureLinear = curvatures["Curvature linear term [1/m^2]"].to_numpy(dtype=float)
+
+    else:
+
+        curvatureLinear = np.zeros(len(curvature))
+
+    ds = positions[1:] - positions[:-1]
+
+    heading = np.zeros(len(positions))
+    heading[1:] = np.cumsum(curvature * ds + 0.5 * curvatureLinear * ds**2)
+
+    return pd.DataFrame(
+        {"Heading [rad]": heading},
+        index=positions
+    )
 
 
 class TestCurvature(unittest.TestCase):
 
-    def testLinearCurvature(self):
+    def test_cvodes_pwc_midpoint_curvature_converges_to_pwl_curvature(self):
+        '''
+        Track with a linearly increasing curvature over 1000 m.
+
+        The result obtained using the piecewise linear curvature model is compared
+        against a piecewise constant midpoint approximation of the curvature with increasing numbers of intervals.
+
+        The piecewise constant approximation should converge to the piecewise linear
+        result for both duration and final velocity.
+
+        CVODES is used as the integrator.
+        '''
+
+        train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
+
+        optsDict = {'integrationMethod': 'CVODES'}
+        opts = OptionsCasadiSolver(optsDict)
+
+        trainModel = train.exportModel()
+        trainIntegrator = TrainIntegrator(trainModel, opts.integrationMethod, opts.integrationOptions.toDict())
+
+        # Scenario
+        time0 = 0
+        velSq0 = 1
+        ds = 1000
+        traction = 0.8 * (train.forceMax / train.mass)
+
+        initialCurvature = 0
+        finalCurvature = 0.005
+
+        maxIntervals = 50
+        relativeTolerance = 1e-4
+
+        # PWL curvature reference
+        out = trainIntegrator.solve(
+            time=time0,
+            velocitySquared=velSq0,
+            ds=ds,
+            traction=traction,
+            pnBrake=0,
+            gradient=0,
+            gradientLinearTerm=0,
+            curvature=initialCurvature,
+            curvatureLinearTerm=(finalCurvature - initialCurvature) / ds,
+            tunnelFactor=0
+        )
+
+        pwlDuration = float(out['time'])
+        pwlVelocity = np.sqrt(float(out['velSquared']))
+
+        # PWC curvature using midpoint rule
+        times = []
+        velocities = []
+        intervalCounts = []
+
+        for numIntervals in range(1, maxIntervals + 1):
+
+            time = time0
+            velSq = velSq0
+
+            for idx in range(numIntervals):
+
+                curvature = (initialCurvature + (idx + 0.5) * (finalCurvature - initialCurvature) / numIntervals)
+
+                out = trainIntegrator.solve(
+                    time=time,
+                    velocitySquared=velSq,
+                    ds=ds / numIntervals,
+                    traction=traction,
+                    pnBrake=0,
+                    gradient=0,
+                    gradientLinearTerm=0,
+                    curvature=curvature,
+                    curvatureLinearTerm=0,
+                    tunnelFactor=0
+                )
+
+                time = out['time']
+                velSq = out['velSquared']
+
+            intervalCounts.append(numIntervals)
+            times.append(float(time))
+            velocities.append(np.sqrt(float(velSq)))
+
+        finalPwcDuration = times[-1]
+        finalPwcVelocity = velocities[-1]
+
+        relativeDurationError = abs(finalPwcDuration - pwlDuration) / pwlDuration
+        relativeVelocityError = abs(finalPwcVelocity - pwlVelocity) / pwlVelocity
+
+        self.assertLess(
+            relativeDurationError,
+            relativeTolerance,
+            msg=(
+                "PWC midpoint approximation did not converge sufficiently to the PWL duration. "
+                f"PWL duration: {pwlDuration:.6f}, "
+                f"PWC duration: {finalPwcDuration:.6f}, "
+                f"relative error: {relativeDurationError:.6e}."
+            )
+        )
+
+        self.assertLess(
+            relativeVelocityError,
+            relativeTolerance,
+            msg=(
+                "PWC midpoint approximation did not converge sufficiently to the PWL final velocity. "
+                f"PWL velocity: {pwlVelocity:.6f}, "
+                f"PWC velocity: {finalPwcVelocity:.6f}, "
+                f"relative error: {relativeVelocityError:.6e}."
+            )
+        )
+
+        if plotDebug:
+            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(16, 10))
+
+            fig.suptitle("CVODES: PWC midpoint curvature approximation compared to PWL curvature", fontsize=14)
+
+            ax1.axhline(pwlDuration, label="pwl")
+            ax1.plot(intervalCounts, times, marker="o", color="orange", label="pwc midpoint")
+            ax1.set_xlabel("Number of intervals of pwc curvature approximation")
+            ax1.set_ylabel("Duration [s]")
+            ax1.grid(True, which="both", linestyle="--", alpha=0.5)
+            ax1.legend(loc="upper right")
+
+            ax2.axhline(pwlVelocity, label="pwl")
+            ax2.plot(intervalCounts, velocities, marker="o", color="orange", label="pwc midpoint")
+            ax2.set_xlabel("Number of intervals of pwc curvature approximation")
+            ax2.set_ylabel("Velocity [m/s]")
+            ax2.grid(True, which="both", linestyle="--", alpha=0.5)
+            ax2.legend(loc="upper right")
+
+            fig.tight_layout()
+            plt.show()
+
+
+    def test_rk_pwc_midpoint_curvature_converges_to_pwl_curvature(self):
+        '''
+        Track with a linearly increasing curvature over 1000 m.
+
+        The result obtained using the piecewise linear curvature model is compared
+        against a piecewise constant midpoint approximation of the curvature with increasing numbers of intervals.
+
+        The piecewise constant approximation should converge to the piecewise linear
+        result for both duration and final velocity.
+
+        RK without time approximation is used as the integrator.
+        RK substeps are increased until convergence
+        '''
+
+        train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
+
+        optsDict = {'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps': 0, 'numSteps': 1}}
+        opts = OptionsCasadiSolver(optsDict)
+
+        trainModel = train.exportModel()
+        trainIntegrator = TrainIntegrator(trainModel, opts.integrationMethod, opts.integrationOptions.toDict())
+
+        # Scenario
+        time0 = 0
+        velSq0 = 1
+        ds = 1000
+        traction = 0.8 * (train.forceMax / train.mass)
+
+        initialCurvature = 0
+        finalCurvature = 0.005
+
+        maxIntervals = 50
+        relativeTolerance = 1e-5
+
+        # PWL gradient reference
+        pwlTimes = []
+        pwlVelocities = []
+        pwlIntervalCounts = []
+
+        for numStep in range(1, maxIntervals + 1):
+
+            optsDict = {'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 0, 'numSteps': numStep}}
+            opts = OptionsCasadiSolver(optsDict)
+            pwlTrainIntegrator = TrainIntegrator(trainModel, opts.integrationMethod, opts.integrationOptions.toDict())
+
+            out = pwlTrainIntegrator.solve(
+                time=time0,
+                velocitySquared=velSq0,
+                ds=ds,
+                traction=traction,
+                pnBrake=0,
+                gradient=0,
+                gradientLinearTerm=0,
+                curvature=initialCurvature,
+                curvatureLinearTerm=(finalCurvature - initialCurvature) / ds,
+                tunnelFactor=0
+            )
+
+            pwlTimes.append(float(out['time']))
+            pwlVelocities.append(np.sqrt(float(out['velSquared'])))
+            pwlIntervalCounts.append(numStep)
+
+        finalPwlDuration = pwlTimes[-1]
+        finalPwlVelocity = pwlVelocities[-1]
+
+        # PWC curvature using midpoint rule
+        pwcTimes = []
+        pwcVelocities = []
+        pwcIntervalCounts = []
+
+        for numIntervals in range(1, maxIntervals + 1):
+
+            time = time0
+            velSq = velSq0
+
+            for idx in range(numIntervals):
+
+                curvature = (initialCurvature + (idx + 0.5) * (finalCurvature - initialCurvature) / numIntervals)
+
+                out = trainIntegrator.solve(
+                    time=time,
+                    velocitySquared=velSq,
+                    ds=ds / numIntervals,
+                    traction=traction,
+                    pnBrake=0,
+                    gradient=0,
+                    gradientLinearTerm=0,
+                    curvature=curvature,
+                    curvatureLinearTerm=0,
+                    tunnelFactor=0
+                )
+
+                time = out['time']
+                velSq = out['velSquared']
+
+            pwcIntervalCounts.append(numIntervals)
+            pwcTimes.append(float(time))
+            pwcVelocities.append(np.sqrt(float(velSq)))
+
+        finalPwcDuration = pwcTimes[-1]
+        finalPwcVelocity = pwcVelocities[-1]
+
+        relativeDurationError = abs(finalPwcDuration - finalPwlDuration) / finalPwlDuration
+        relativeVelocityError = abs(finalPwcVelocity - finalPwlVelocity) / finalPwlVelocity
+
+        self.assertLess(
+            relativeDurationError,
+            relativeTolerance,
+            msg=(
+                "PWC midpoint approximation did not converge sufficiently to the PWL duration. "
+                f"PWL duration: {finalPwlDuration:.6f}, "
+                f"PWC duration: {finalPwcDuration:.6f}, "
+                f"relative error: {relativeDurationError:.6e}."
+            )
+        )
+
+        self.assertLess(
+            relativeVelocityError,
+            relativeTolerance,
+            msg=(
+                "PWC midpoint approximation did not converge sufficiently to the PWL final velocity. "
+                f"PWL velocity: {finalPwlVelocity:.6f}, "
+                f"PWC velocity: {finalPwcVelocity:.6f}, "
+                f"relative error: {relativeVelocityError:.6e}."
+            )
+        )
+
+        if plotDebug:
+            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(16, 10))
+
+            fig.suptitle("Explicit RK: PWL curvature with RK substeps vs PWC midpoint approximation", fontsize=14)
+
+            ax1.plot(pwlIntervalCounts, pwlTimes, marker="o", label="PWL: RK substeps")
+            ax1.plot(pwcIntervalCounts, pwcTimes, marker="o", color="orange", label="PWC: intervals")
+            ax1.set_xlabel("Refinement level: PWC intervals and RK substeps")
+            ax1.set_ylabel("Duration [s]")
+            ax1.grid(True, which="both", linestyle="--", alpha=0.5)
+            ax1.legend(loc="upper right")
+
+            ax2.plot(pwlIntervalCounts, pwlVelocities, marker="o", label="PWL: RK substeps")
+            ax2.plot(pwcIntervalCounts, pwcVelocities, marker="o", color="orange", label="PWC: intervals")
+            ax2.set_xlabel("Refinement level: PWC intervals and RK substeps")
+            ax2.set_ylabel("Velocity [m/s]")
+            ax2.grid(True, which="both", linestyle="--", alpha=0.5)
+            ax2.legend(loc="upper right")
+
+            fig.tight_layout()
+            plt.show()
+
+
+    def test_rk_time_approx_pwc_midpoint_curvature_converges_to_pwl_curvature(self):
+        '''
+        Track with a linearly increasing curvature over 1000 m.
+
+        The result obtained using the piecewise linear curvature model is compared
+        against a piecewise constant midpoint approximation of the curvature.
+
+        The piecewise constant approximation should converge to the piecewise linear
+        result for both duration and final velocity.
+
+        RK with time approximation is used as the integrator.
+        RK uses 50 substeps.
+        Time approx steps are increased until convergence
+        '''
+
+        train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
+        trainModel = train.exportModel()
+
+        # Scenario
+        time0 = 0
+        velSq0 = 1
+        ds = 1000
+        traction = 0.8 * (train.forceMax / train.mass)
+
+        initialCurvature = 0
+        finalCurvature = 0.005
+
+        numIntervals = 50
+        timeApproxSteps = 30
+        relativeTolerance = 1e-3
+
+        # PWL gradient reference
+        pwlTimes = []
+        pwlVelocities = []
+        timeApproxStepCounts  = []
+
+        for timeSteps in range(1, timeApproxSteps + 1):
+
+            optsDict = {'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': timeSteps, 'numSteps': 50}}
+            opts = OptionsCasadiSolver(optsDict)
+            trainIntegrator = TrainIntegrator(trainModel, opts.integrationMethod, opts.integrationOptions.toDict())
+
+            out = trainIntegrator.solve(
+                time=time0,
+                velocitySquared=velSq0,
+                ds=ds,
+                traction=traction,
+                pnBrake=0,
+                gradient=0,
+                gradientLinearTerm=0,
+                curvature=initialCurvature,
+                curvatureLinearTerm=(finalCurvature - initialCurvature) / ds,
+                tunnelFactor=0
+            )
+
+            pwlTimes.append(float(out['time']))
+            pwlVelocities.append(np.sqrt(float(out['velSquared'])))
+            timeApproxStepCounts.append(timeSteps)
+
+        finalPwlDuration = pwlTimes[-1]
+        finalPwlVelocity = pwlVelocities[-1]
+
+        # PWC curvature using midpoint rule
+        pwcTimes = []
+        pwcVelocities = []
+
+        for timeSteps in range(1, timeApproxSteps + 1):
+
+            optsDict = {'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': timeSteps, 'numSteps': 50}}
+            opts = OptionsCasadiSolver(optsDict)
+            trainIntegrator = TrainIntegrator(trainModel, opts.integrationMethod, opts.integrationOptions.toDict())
+
+            time = time0
+            velSq = velSq0
+
+            for idx in range(numIntervals):
+
+                curvature = (initialCurvature + (idx + 0.5) * (finalCurvature - initialCurvature) / numIntervals)
+
+                out = trainIntegrator.solve(
+                    time=time,
+                    velocitySquared=velSq,
+                    ds=ds / numIntervals,
+                    traction=traction,
+                    pnBrake=0,
+                    gradient=0,
+                    gradientLinearTerm=0,
+                    curvature=curvature,
+                    curvatureLinearTerm=0,
+                    tunnelFactor=0
+                )
+
+                time = out['time']
+                velSq = out['velSquared']
+
+            pwcTimes.append(float(time))
+            pwcVelocities.append(np.sqrt(float(velSq)))
+
+        finalPwcDuration = pwcTimes[-1]
+        finalPwcVelocity = pwcVelocities[-1]
+
+        relativeDurationError = abs(finalPwcDuration - finalPwlDuration) / finalPwlDuration
+        relativeVelocityError = abs(finalPwcVelocity - finalPwlVelocity) / finalPwlVelocity
+
+        self.assertLess(
+            relativeDurationError,
+            relativeTolerance,
+            msg=(
+                "PWC midpoint approximation did not converge sufficiently to the PWL duration. "
+                f"PWL duration: {finalPwlDuration:.6f}, "
+                f"PWC duration: {finalPwcDuration:.6f}, "
+                f"relative error: {relativeDurationError:.6e}."
+            )
+        )
+
+        self.assertLess(
+            relativeVelocityError,
+            relativeTolerance,
+            msg=(
+                "PWC midpoint approximation did not converge sufficiently to the PWL final velocity. "
+                f"PWL velocity: {finalPwlVelocity:.6f}, "
+                f"PWC velocity: {finalPwcVelocity:.6f}, "
+                f"relative error: {relativeVelocityError:.6e}."
+            )
+        )
+
+        if plotDebug:
+            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(16, 10))
+
+            fig.suptitle("RK with Time Approx: PWC midpoint curvature approximation compared to PWL curvature", fontsize=14)
+
+            ax1.plot(timeApproxStepCounts , pwlTimes, marker="o", label="pwl")
+            ax1.plot(timeApproxStepCounts, pwcTimes, marker="o", color="orange", label="pwc midpoint")
+            ax1.set_xlabel("Number of time approx steps")
+            ax1.set_ylabel("Duration [s]")
+            ax1.grid(True, which="both", linestyle="--", alpha=0.5)
+            ax1.legend(loc="upper right")
+
+            ax2.plot(timeApproxStepCounts , pwlVelocities, marker="o", label="pwl")
+            ax2.plot(timeApproxStepCounts, pwcVelocities, marker="o", color="orange", label="pwc midpoint")
+            ax2.set_xlabel("Number of time approx steps")
+            ax2.set_ylabel("Velocity [m/s]")
+            ax2.grid(True, which="both", linestyle="--", alpha=0.5)
+            ax2.legend(loc="upper right")
+            ax2.ticklabel_format(axis="y", style="plain", useOffset=False)
+
+            fig.tight_layout()
+            plt.show()
+
+
+    def test_train_length_dependent_curvature_preserves_target_heading(self):
+        '''
+        Compare the final heading of the original length-independent curvature profile
+        with the train-length-dependent piecewise linear curvature profile.
+
+        Both profiles should start from the same heading and end at the same target heading.
+        '''
+
+        headingTolerance  = 1e-6
+
+        trainLength = 800   # [m]
+        track = Track(config={'id': 'CH_StGallen_Wil'}, pathJSON='tracks')
+
+        # track needs to be straight at least train length meters before the end of the track
+        track.curvatures = track.curvatures[track.curvatures.index < track.length - trainLength]
+        track.curvatures.loc[track.length - trainLength] = {"Curvature [1/m]": 0.0}
+
+        df_heading_indep = computeHeadingFromCurvature(track.curvatures, track.length)
+
+        train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
+        train.length = trainLength
+        track.updateTrainLengthDependentValues(train)
+
+        df_heading_dep = computeHeadingFromCurvature(track.curvatures, track.length)
+
+        self.assertAlmostEqual(
+            df_heading_indep["Heading [rad]"].iloc[0],
+            df_heading_dep["Heading [rad]"].iloc[0],
+            delta=headingTolerance,
+            msg="Length-independent and train-length-dependent curvature profiles should start with the same heading."
+        )
+
+        self.assertAlmostEqual(
+            df_heading_indep["Heading [rad]"].iloc[-1],
+            df_heading_dep["Heading [rad]"].iloc[-1],
+            delta=headingTolerance,
+            msg=(
+                "Length-independent and train-length-dependent curvature profiles "
+                "should end with the same heading. "
+                f"Length-independent final heading: {df_heading_indep['Heading [rad]'].iloc[-1]:.12f} rad, "
+                f"train-length-dependent final heading: {df_heading_dep['Heading [rad]'].iloc[-1]:.12f} rad."
+            )
+        )
+
+        if plotDebug:
+
+            fig, ax = plt.subplots(figsize=(16, 8))
+
+            ax.plot(df_heading_indep.index.values / 1000, df_heading_indep["Heading [rad]"].to_numpy(), label="length-independent heading")
+            ax.plot(df_heading_dep.index.values / 1000, df_heading_dep["Heading [rad]"].to_numpy(), label="length-dependent heading")
+            ax.set_title("Heading  Comparison")
+            ax.set_xlabel("Position [km]")
+            ax.set_ylabel("Heading [rad]")
+            ax.grid(True, which="both", linestyle="--", alpha=0.5)
+            ax.legend(loc="upper right")
+            ax.set_xlim(0, track.length / 1000)
+            ax.figure.tight_layout()
+
+            plt.show()
+
+
+    def test_pwc_curvature_approximation_matches_pwl_curvature_energy(self):
         '''
         Track with right turn from 1000 m to 2000 m and a left turn from 3000 m to 4000 m.
 
@@ -17,31 +546,33 @@ def testLinearCurvature(self):
         linear curvatures or equivalent piecewise constant curvatures.
         '''
 
-        startPosition = 0 # [m]
-        endPosition = 5000 # [m]
-        duration = 5000/(60/3.6) # [s]
+        startPosition = 0  # [m]
+        endPosition = 5000  # [m]
+        duration = 5000/(60/3.6)  # [s]
 
-        energyRelativeTolerance = 0.004
+        energyRelativeTolerance = 1e-4
         numIntervals = 100
 
-        train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='trains')
+        train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
         train.length = 600
 
         track = Track(config={'id': 'test_two_radii'}, pathJSON='tracks')
         track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
         track.updateTrainLengthDependentValues(train)
 
-        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True}
+        # PWL Curvatures
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'CVODES'}
         solver = casadiSolver(train, track, opts)
-        df1, stats1 = solver.solve(duration)
+        pwl_df, pwl_stats = solver.solve(duration)
 
-        energyConsumptionWithLinearTerms = stats1['Cost']
+        energyConsumptionWithLinearTerms = pwl_stats['Cost']
 
-        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True, 'pwcLengthDependentTrackAttributes': True}
+        # PWC Curvatures
+        opts = {'numIntervals': numIntervals, 'integrationMethod': 'CVODES', 'pwcLengthDependentTrackAttributes': True}
         solver = casadiSolver(train, track, opts)
-        df2, stats2 = solver.solve(duration)
+        pwc_df, pwc_stats = solver.solve(duration)
 
-        energyConsumptionWithPwcTerms= stats2['Cost']
+        energyConsumptionWithPwcTerms= pwc_stats['Cost']
 
         relativeEnergyDifference = (abs(energyConsumptionWithLinearTerms - energyConsumptionWithPwcTerms) / energyConsumptionWithLinearTerms)
 
@@ -57,15 +588,12 @@ def testLinearCurvature(self):
             )
         )
 
-
-        plotDebug = True
-
         if plotDebug:
 
             fig, ax = plt.subplots(figsize=(16, 8))
 
-            ax.plot(df1["Position [m]"] / 1000, df1["Curvature [1/m]"], label="pwl curvatures")
-            ax.step(df2["Position [m]"] / 1000, df2["Curvature [1/m]"], "--", where="post", label="pwc curvatures")
+            ax.plot(pwl_df["Position [m]"] / 1000, pwl_df["Curvature [1/m]"], label="pwl curvatures")
+            ax.step(pwc_df["Position [m]"] / 1000, pwc_df["Curvature [1/m]"], "--", where="post", label="pwc curvatures")
             ax.set_title("Curvatures")
             ax.set_xlabel("Position [km]")
             ax.set_ylabel("Curvature [1/m]")
@@ -76,93 +604,52 @@ def testLinearCurvature(self):
 
             plt.show()
 
-    def testAllIntegratorTypesWork(self):
-        '''
-        Verify that all supported integration methods produce consistent results
-        for the same train, track, and optimization setup.
 
-        The test compares RK, IRK, and CVODES, including the approximate time
-        integration option for RK and IRK. The resulting energy costs should only
-        differ by a small relative tolerance.
+    def test_short_train_length_has_negligible_effect_on_curvature_energy_consumption(self):
         '''
+        Use an artificially short train on a real track profile.
 
-        startPosition = 0  # [m]
-        endPosition = 5000  # [m]
-        duration = 5000 / (60 / 3.6)  # [s]
-
-        tol = 0.02
-        numIntervals = 100
-
-        train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='trains')
-        train.length = 600
-
-        track = Track(config={'id': 'test_two_radii'}, pathJSON='tracks')
-        track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
-        track.updateTrainLengthDependentValues(train)
-
-        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1},'energyOptimal': True}
-        solver = casadiSolver(train, track, opts)
-        df, stats = solver.solve(duration)
-
-        energy_RK_Approx = stats['Cost']
+        For a very small train length, the train-length-dependent curvature profile
+        should be almost identical to the original curvature profile. Therefore, the
+        energy consumption should remain within a small relative tolerance.
+        '''
 
-        opts = {'numIntervals': numIntervals, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 0},'energyOptimal': True}
-        solver = casadiSolver(train, track, opts)
-        df, stats = solver.solve(duration)
+        relativeTolerance = 1e-2
+        trainLength = 10  # [m]
+        train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
+        train.length = trainLength
 
-        energy_RK = stats['Cost']
+        track = Track(config={'id': 'CH_StGallen_Wil'}, pathJSON='tracks')
+        track.gradients = track.gradients.iloc[[0]]
+        track.gradients["Gradient [permil]"].iloc[0] = 0
 
-        opts = {'numIntervals': numIntervals, 'integrationMethod': 'IRK', 'integrationOptions': {'numApproxSteps': 1},'energyOptimal': True}
-        solver = casadiSolver(train, track, opts)
-        df, stats = solver.solve(duration)
+        # track needs to be straight at least train length meters before the end of the track
+        track.curvatures = track.curvatures[track.curvatures.index < track.length - trainLength]
+        track.curvatures.loc[track.length - trainLength] = {"Curvature [1/m]": 0.0}
 
-        energy_IRK_Approx = stats['Cost']
+        duration = track.length / (80/3.6)
 
-        opts = {'numIntervals': numIntervals, 'integrationMethod': 'IRK', 'integrationOptions': {'numApproxSteps': 0},'energyOptimal': True}
+        # train-length-independent
+        opts = {'numIntervals':600, 'integrationMethod':'RK', 'integrationOptions':{'numApproxSteps':0}, 'energyOptimal':True}
         solver = casadiSolver(train, track, opts)
-        df, stats = solver.solve(duration)
+        indep_df, indep_stats = solver.solve(duration)
 
-        energy_IRK = stats['Cost']
-
-        opts = {'numIntervals': numIntervals, 'integrationMethod': 'CVODES', 'energyOptimal': True}
+        track.updateTrainLengthDependentValues(train)
         solver = casadiSolver(train, track, opts)
-        df, stats = solver.solve(duration)
-
-        energy_CVODES = stats['Cost']
+        dep_df, dep_stats = solver.solve(duration)
 
-        relDiff_RKApprox_IRKApprox = abs(energy_RK_Approx - energy_IRK_Approx) / energy_RK_Approx
-        relDiff_RKApprox_CVODES = abs(energy_RK_Approx - energy_CVODES) / energy_RK_Approx
-        relDiff_RK_IRK = abs(energy_RK - energy_IRK) / energy_RK
-
-        self.assertLess(
-            relDiff_RKApprox_IRKApprox,
-            tol,
-            msg=(
-                "RK and IRK with numApproxSteps=1 should give similar costs. "
-                f"RK approx: {energy_RK_Approx:.6f}, "
-                f"IRK approx: {energy_IRK_Approx:.6f}, "
-                f"relative difference: {relDiff_RKApprox_IRKApprox:.6f}."
-            )
-        )
+        energyConsumptionIndependentOfTrainLength = indep_stats['Cost']
+        energyConsumptionDependentOfTrainLength = dep_stats['Cost']
 
-        self.assertLess(
-            relDiff_RKApprox_CVODES,
-            tol,
-            msg=(
-                "RK with numApproxSteps=1 and CVODES should give similar costs. "
-                f"RK approx: {energy_RK_Approx:.6f}, "
-                f"CVODES: {energy_CVODES:.6f}, "
-                f"relative difference: {relDiff_RKApprox_CVODES:.6f}."
-            )
-        )
+        relativeDifference = (abs(energyConsumptionDependentOfTrainLength - energyConsumptionIndependentOfTrainLength) / energyConsumptionIndependentOfTrainLength)
 
         self.assertLess(
-            relDiff_RK_IRK,
-            tol,
+            relativeDifference,
+            relativeTolerance,
             msg=(
-                "RK and IRK with numApproxSteps=0 should give similar costs. "
-                f"RK: {energy_RK:.6f}, "
-                f"IRK: {energy_IRK:.6f}, "
-                f"relative difference: {relDiff_RK_IRK:.6f}."
+                "Energy consumption with and without train-length-dependent curvature should be roughly equal. "
+                f"Independent: {energyConsumptionIndependentOfTrainLength:.6f}, "
+                f"dependent: {energyConsumptionDependentOfTrainLength:.6f}, "
+                f"relative difference: {relativeDifference:.6e}."
             )
         )
\ No newline at end of file
diff --git a/unitTests/trainLengthDependentTrackAttributes/gradient.py b/unitTests/trainLengthDependentTrackAttributes/gradient.py
index 4bd7c78..d1e0774 100644
--- a/unitTests/trainLengthDependentTrackAttributes/gradient.py
+++ b/unitTests/trainLengthDependentTrackAttributes/gradient.py
@@ -1,5 +1,4 @@
 import unittest
-from time import perf_counter_ns
 
 import numpy as np
 from matplotlib import pyplot as plt
@@ -9,9 +8,12 @@
 from train import Train, TrainIntegrator
 
 
+plotDebug = False
+
+
 class TestGradient(unittest.TestCase):
 
-    def test_integrator_CVODES(self):
+    def test_cvodes_pwc_midpoint_gradient_converges_to_pwl_gradient(self):
         '''
         Track with a linearly increasing gradient over 1000 m.
 
@@ -43,7 +45,6 @@ def test_integrator_CVODES(self):
 
         maxIntervals = 50
         relativeTolerance = 1e-3
-        plotDebug = True
 
         # PWL gradient reference
         out = trainIntegrator.solve(
@@ -149,7 +150,7 @@ def test_integrator_CVODES(self):
             plt.show()
 
 
-    def test_integrator_RK(self):
+    def test_rk_pwc_midpoint_gradient_converges_to_pwl_gradient(self):
         '''
         Track with a linearly increasing gradient over 1000 m.
 
@@ -182,7 +183,6 @@ def test_integrator_RK(self):
 
         maxIntervals = 50
         relativeTolerance = 1e-3
-        plotDebug = True
 
         # PWL gradient reference
         pwlTimes = []
@@ -302,7 +302,7 @@ def test_integrator_RK(self):
             plt.show()
 
 
-    def test_integrator_RK_with_Time_Approx(self):
+    def test_rk_time_approx_pwc_midpoint_gradient_converges_to_pwl_gradient(self):
         '''
         Track with a linearly increasing gradient over 1000 m.
 
@@ -332,7 +332,6 @@ def test_integrator_RK_with_Time_Approx(self):
         numIntervals = 50
         timeApproxSteps = 30
         relativeTolerance = 1e-3
-        plotDebug = True
 
         # PWL gradient reference
         pwlTimes = []
@@ -455,7 +454,7 @@ def test_integrator_RK_with_Time_Approx(self):
             plt.show()
 
 
-    def testPWLProfile(self):
+    def test_train_length_dependent_gradient_preserves_target_altitude(self):
         '''
         Compare the final altitude of the original length-independent gradient profile
         with the train-length-dependent piecewise linear profile.
@@ -463,7 +462,6 @@ def testPWLProfile(self):
         Both profiles should start from the same altitude and end at the same target altitude.
         '''
 
-        plotDebug = True
         altitudeTolerance = 1e-6
 
         trainLength = 800   # [m]
@@ -471,7 +469,7 @@ def testPWLProfile(self):
 
         # track needs to be flat at least train length meters before the end of the track
         track.gradients = track.gradients[track.gradients.index < track.length - trainLength]
-        track.gradients.loc[track.length - trainLength] = {"Gradient [permil]": 0.0, "Gradient linear term [permil/m]": 0.0}
+        track.gradients.loc[track.length - trainLength] = {"Gradient [permil]": 0.0}
 
         df_alt = computeAltitude(track.gradients, track.length)
 
@@ -529,7 +527,7 @@ def testPWLProfile(self):
             plt.show()
 
 
-    def testLinearGradient(self):
+    def test_pwc_gradient_approximation_matches_pwl_gradient_energy(self):
         '''
         Track with 20 permil increase from 1000 m to 2000 m and 20 permil
         decrease from 3000 m to 4000 m.
@@ -540,9 +538,9 @@ def testLinearGradient(self):
         Altitude should be 0 m at the end.
         '''
 
-        startPosition = 0 # [m]
-        endPosition = 5000 # [m]
-        duration = 5000/(60/3.6) # [s]
+        startPosition = 0  # [m]
+        endPosition = 5000  # [m]
+        duration = 5000/(60/3.6)  # [s]
 
         altitudeTolerance = 1e-4
         energyRelativeTolerance = 1e-3
@@ -597,8 +595,6 @@ def testLinearGradient(self):
             )
         )
 
-        plotDebug = True
-
         if plotDebug:
 
             fig, ax = plt.subplots(figsize=(16, 8))
@@ -617,7 +613,7 @@ def testLinearGradient(self):
             plt.show()
 
 
-    def testSmallTrainLengthNotAffectingEnergyConsumption(self):
+    def test_short_train_length_has_negligible_effect_on_energy_consumption(self):
         '''
         Use an artificially short train on a real track profile.
 
@@ -632,10 +628,12 @@ def testSmallTrainLengthNotAffectingEnergyConsumption(self):
         train.length = trainLength
 
         track = Track(config={'id': 'CH_StGallen_Wil'}, pathJSON='tracks')
+        track.curvatures = track.curvatures.iloc[[0]]
+        track.curvatures["Curvature [1/m]"].iloc[0] = 0
 
         # track needs to be flat at least train length meters before the end of the track
         track.gradients = track.gradients[track.gradients.index < track.length - trainLength]
-        track.gradients.loc[track.length - trainLength] = {"Gradient [permil]": 0.0, "Gradient linear term [permil/m]": 0.0}
+        track.gradients.loc[track.length - trainLength] = {"Gradient [permil]": 0.0}
 
         duration = track.length / (80/3.6)
 

From 5b58c065d1b21c2689056a216cdf91a6efcb729b Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Wed, 3 Jun 2026 12:09:44 +0200
Subject: [PATCH 26/31] bug fixes

---
 mseetc/train.py                               |  2 +-
 simulations/sim_launcher.py                   | 17 +++++---
 tracks/swisstopo/analyzeTracks.py             |  6 +--
 trains/CH_Stadler_FLIRT_TPF.json              |  6 +--
 .../curvatureResistance.py                    |  1 -
 unitTests/integrators/integrators.py          |  6 +--
 .../curvature.py                              |  6 +--
 .../gradient.py                               |  6 +--
 .../speedLimit.py                             | 43 ++++++++++---------
 .../tunnelResistance/tunnelResistance.py      | 40 ++++++++---------
 10 files changed, 70 insertions(+), 63 deletions(-)

diff --git a/mseetc/train.py b/mseetc/train.py
index 73d08f8..7176ae5 100644
--- a/mseetc/train.py
+++ b/mseetc/train.py
@@ -117,7 +117,7 @@ def __init__(self, config, pathJSON=Path(__file__).parent.parent / 'trains') ->
             tunnel_data = data["tunnel resistance"] # additive aerodynamic tunnel drag as dict per tunnel cross section
 
             cross_section_unit = tunnel_data["units"]["cross section"] # tunnel cross section [m^2]
-            resistance_unit = tunnel_data["units"]["resistance (or similar)"] # resistance coefficient [kg/m]
+            resistance_unit = tunnel_data["units"]["coefficient"] # resistance coefficient [kg/m]
 
             self.tunnelCoefficients = {
                 convertUnit(cross_section, cross_section_unit): convertUnit(
diff --git a/simulations/sim_launcher.py b/simulations/sim_launcher.py
index c08acd5..515caaa 100644
--- a/simulations/sim_launcher.py
+++ b/simulations/sim_launcher.py
@@ -1,31 +1,38 @@
-from efficiency import totalLossesFunction
-from ocp import casadiSolver
+from mseetc.efficiency import totalLossesFunction
+from mseetc.ocp import casadiSolver
+
 
 def get_power_loss_function(train, mode="perfect",* ,auxiliaries: float = 27_000, eta_gear: float = 0.96):
 
     if mode == "perfect":
+
         return lambda f, v: 0
 
     elif mode == "static":
+
         return lambda f, v: (f>0)*f*v*(1-train.etaTraction)/train.etaTraction - (f<0)*f*v*(1-train.etaRgBrake)
 
     elif mode == "dynamic":
+
         return totalLossesFunction(train, auxiliaries=auxiliaries, etaGear=eta_gear)
 
     else:
+
         raise ValueError("mode must be one of: 'perfect', 'static', 'dynamic'")
 
+
+
 if __name__ == '__main__':
 
-    from train import Train
-    from track import Track
+    from mseetc.train import Train
+    from mseetc.track import Track
 
     # Timetable
     startPosition = 0       # [m]
     endPosition = 20000     # [m]
     duration = 60*20        # [s]
 
-    train = Train(config={'id':'Flirt_Tpf'}, pathJSON='../trains')
+    train = Train(config={'id':'CH_Stadler_FLIRT_TPF'}, pathJSON='../trains')
     train.forceMinPn = 0
     train.withPnBrake = False
     train.powerLosses = get_power_loss_function(train, "static")
diff --git a/tracks/swisstopo/analyzeTracks.py b/tracks/swisstopo/analyzeTracks.py
index 70a8cf1..250bea5 100644
--- a/tracks/swisstopo/analyzeTracks.py
+++ b/tracks/swisstopo/analyzeTracks.py
@@ -1,10 +1,10 @@
 import numpy as np
 from matplotlib import pyplot as plt
 
-from ocp import casadiSolver
+from mseetc.ocp import casadiSolver
 from simulations.sim_launcher import get_power_loss_function
-from track import Track
-from train import Train
+from mseetc.track import Track
+from mseetc.train import Train
 
 
 if __name__ == '__main__':
diff --git a/trains/CH_Stadler_FLIRT_TPF.json b/trains/CH_Stadler_FLIRT_TPF.json
index a0e0ff9..9443e2b 100644
--- a/trains/CH_Stadler_FLIRT_TPF.json
+++ b/trains/CH_Stadler_FLIRT_TPF.json
@@ -80,16 +80,16 @@
     "tunnel resistance": {
         "units": {
             "cross section": "m^2",
-            "resistance (or similar)": "kg/m"
+            "coefficient": "kN/(km/h)^2"
         },
         "values": [
             [
                 24,
-                15.16
+                0.0011698
             ],
             [
                 40,
-                7.23
+                0.0005579
             ]
         ]
     }
diff --git a/unitTests/curvatureResistance/curvatureResistance.py b/unitTests/curvatureResistance/curvatureResistance.py
index cf6bbac..e7e8ec0 100644
--- a/unitTests/curvatureResistance/curvatureResistance.py
+++ b/unitTests/curvatureResistance/curvatureResistance.py
@@ -1,6 +1,5 @@
 import copy
 import numpy as np
-import sys
 import unittest
 
 from mseetc.efficiency import totalLossesFunction
diff --git a/unitTests/integrators/integrators.py b/unitTests/integrators/integrators.py
index 649802a..062aa64 100644
--- a/unitTests/integrators/integrators.py
+++ b/unitTests/integrators/integrators.py
@@ -1,8 +1,8 @@
 import unittest
 
-from ocp import casadiSolver
-from track import Track
-from train import Train
+from mseetc.ocp import casadiSolver
+from mseetc.track import Track
+from mseetc.train import Train
 
 
 class TestGradient(unittest.TestCase):
diff --git a/unitTests/trainLengthDependentTrackAttributes/curvature.py b/unitTests/trainLengthDependentTrackAttributes/curvature.py
index 932d2b0..1b3bba7 100644
--- a/unitTests/trainLengthDependentTrackAttributes/curvature.py
+++ b/unitTests/trainLengthDependentTrackAttributes/curvature.py
@@ -4,9 +4,9 @@
 import pandas as pd
 from matplotlib import pyplot as plt
 
-from ocp import casadiSolver, OptionsCasadiSolver
-from track import Track
-from train import Train, TrainIntegrator
+from mseetc.ocp import casadiSolver, OptionsCasadiSolver
+from mseetc.track import Track
+from mseetc.train import Train, TrainIntegrator
 
 
 plotDebug = False
diff --git a/unitTests/trainLengthDependentTrackAttributes/gradient.py b/unitTests/trainLengthDependentTrackAttributes/gradient.py
index d1e0774..6fd5018 100644
--- a/unitTests/trainLengthDependentTrackAttributes/gradient.py
+++ b/unitTests/trainLengthDependentTrackAttributes/gradient.py
@@ -3,9 +3,9 @@
 import numpy as np
 from matplotlib import pyplot as plt
 
-from ocp import casadiSolver, OptionsCasadiSolver
-from track import Track, computeAltitude
-from train import Train, TrainIntegrator
+from mseetc.ocp import casadiSolver, OptionsCasadiSolver
+from mseetc.track import Track, computeAltitude
+from mseetc.train import Train, TrainIntegrator
 
 
 plotDebug = False
diff --git a/unitTests/trainLengthDependentTrackAttributes/speedLimit.py b/unitTests/trainLengthDependentTrackAttributes/speedLimit.py
index d5a9a1a..8b99ffe 100644
--- a/unitTests/trainLengthDependentTrackAttributes/speedLimit.py
+++ b/unitTests/trainLengthDependentTrackAttributes/speedLimit.py
@@ -1,13 +1,13 @@
 import unittest
 
-from ocp import casadiSolver
-from track import Track
-from train import Train
+from mseetc.ocp import casadiSolver
+from mseetc.track import Track
+from mseetc.train import Train
 
 
 class TestSpeedLimit(unittest.TestCase):
 
-    def testTrainDoesNotAccelerateToEarly(self):
+    def test_train_length_dependent_speed_limit_delays_acceleration(self):
         '''
         Speed limit increases from 22 m/s to 40 m/s at position 1000 m.
 
@@ -16,46 +16,47 @@ def testTrainDoesNotAccelerateToEarly(self):
         22 m/s after the whole train has passed the speed-increase position.
         '''
 
-        startPosition = 0 # [m]
-        endPosition = 12000 # [m]
-        duration = 12000/(115/3.6) # [s]
+        startPosition = 0  # [m]
+        endPosition = 12000  # [m]
+        duration = endPosition/(115/3.6)  # [s]
+        speedIncreasePosition = 1000  # [m]
+        speedLimitBeforeIncrease = 22  # [m/s]
+        speedTolerance = 0.001  # [m/s]
 
-        train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='trains')
+        train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
 
         track = Track(config={'id': 'test_one_speed_increase'}, pathJSON='tracks')
         track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
 
         opts = {'numIntervals': 300, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True}
         solver = casadiSolver(train, track, opts)
-        df1, stats1 = solver.solve(duration)
+        dfWithoutTrainLength, statsWithoutTrainLength = solver.solve(duration)
 
-        speedAfterSpeedIncrease1 = df1[df1['Position [m]'] > 1000].iloc[0]['Velocity [m/s]']
+        speedWithoutTrainLengthAfterIncrease = dfWithoutTrainLength[dfWithoutTrainLength['Position [m]'] > speedIncreasePosition].iloc[0]['Velocity [m/s]']
 
         track.updateTrainLengthDependentValues(train)
 
         opts = {'numIntervals': 300, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True}
         solver = casadiSolver(train, track, opts)
-        df2, stats2 = solver.solve(duration)
+        dfWithTrainLength, statsWithTrainLength = solver.solve(duration)
 
-        speedAfterSpeedIncrease2 = df2[df2['Position [m]'] > 1000].iloc[0]['Velocity [m/s]']
-        speedAfterSpeedIncrease3 = df2[df2['Position [m]'] > (1000+train.length)].iloc[0]['Velocity [m/s]']
-
-        epsilon = 0.001
+        speedWithTrainLengthAfterIncrease = dfWithTrainLength[dfWithTrainLength['Position [m]'] > speedIncreasePosition].iloc[0]['Velocity [m/s]']
+        speedWithTrainLengthAfterTrainPassedIncrease = dfWithTrainLength[dfWithTrainLength['Position [m]'] > (speedIncreasePosition+train.length)].iloc[0]['Velocity [m/s]']
 
         self.assertGreater(
-            speedAfterSpeedIncrease1,
-            22,
+            speedWithoutTrainLengthAfterIncrease,
+            speedLimitBeforeIncrease,
             msg="Without train-length-dependent values, the train should accelerate immediately after the speed limit increase."
         )
 
         self.assertLessEqual(
-            speedAfterSpeedIncrease2,
-            22 + epsilon,
+            speedWithTrainLengthAfterIncrease,
+            speedLimitBeforeIncrease + speedTolerance,
             msg="With train-length-dependent values, the train should not accelerate before the whole train has passed the speed limit increase."
         )
 
         self.assertGreater(
-            speedAfterSpeedIncrease3,
-            22,
+            speedWithTrainLengthAfterTrainPassedIncrease,
+            speedLimitBeforeIncrease,
             msg="With train-length-dependent values, the train should accelerate after the whole train has passed the speed limit increase."
         )
\ No newline at end of file
diff --git a/unitTests/tunnelResistance/tunnelResistance.py b/unitTests/tunnelResistance/tunnelResistance.py
index 83df459..e3f600f 100644
--- a/unitTests/tunnelResistance/tunnelResistance.py
+++ b/unitTests/tunnelResistance/tunnelResistance.py
@@ -1,42 +1,42 @@
 import unittest
 
-from ocp import casadiSolver
-from track import Track
-from train import Train
+from mseetc.ocp import casadiSolver
+from mseetc.track import Track
+from mseetc.train import Train
 
 
 class TestTunnelResistance(unittest.TestCase):
 
-    def testHigherEnergyConsumptionInTunnels(self):
+    def test_tunnel_resistance_increases_energy_consumption(self):
         '''
         26 km long small tunnel with cross section of 24 m^2 on a track of 28 km results in significant higher energy consumption.
         '''
 
-        startPosition = 0 # [m]
-        endPosition = 28000 # [m]
-        duration = 28000/(145/3.6) # [s]
+        startPosition = 0  # [m]
+        endPosition = 28000  # [m]
+        duration = 28000/(145/3.6)  # [s]
 
-        train = Train(config={'id': 'Flirt_Tpf'}, pathJSON='trains')
+        minEnergyRatio = 1.5
+
+        train = Train(config={'id': 'CH_Stadler_Flirt_TPF'}, pathJSON='trains')
 
-        track = Track(config={'id': 'test_flat_no_tunnel'}, pathJSON='tracks')
-        track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
+        trackWithoutTunnel = Track(config={'id': 'test_flat_no_tunnel'}, pathJSON='tracks')
+        trackWithoutTunnel.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
 
         opts = {'numIntervals': 300, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True}
-        solver = casadiSolver(train, track, opts)
-        df1, stats1 = solver.solve(duration)
+        solver = casadiSolver(train, trackWithoutTunnel, opts)
+        dfWithoutTunnel, statsWithoutTunnel = solver.solve(duration)
 
-        energyConsumptionWithoutTunnel = stats1['Cost']
+        energyConsumptionWithoutTunnel = statsWithoutTunnel['Cost']
 
-        track = Track(config={'id': 'test_flat_with_tunnel'}, pathJSON='tracks')
-        track.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
+        trackWithTunnel = Track(config={'id': 'test_flat_with_tunnel'}, pathJSON='tracks')
+        trackWithTunnel.updateLimits(positionStart=startPosition, positionEnd=endPosition, unit='m')
 
         opts = {'numIntervals': 300, 'integrationMethod': 'RK', 'integrationOptions': {'numApproxSteps': 1}, 'energyOptimal': True}
-        solver = casadiSolver(train, track, opts)
-        df2, stats2 = solver.solve(duration)
+        solver = casadiSolver(train, trackWithTunnel, opts)
+        dfWithTunnel, statsWithTunnel = solver.solve(duration)
 
-        energyConsumptionWithTunnel = stats2['Cost']
-
-        minEnergyRatio = 1.5
+        energyConsumptionWithTunnel = statsWithTunnel['Cost']
 
         self.assertGreater(
             energyConsumptionWithTunnel,

From 914d72ad6b4844a59805dc2d12ca70424583efc9 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Thu, 4 Jun 2026 16:19:18 +0200
Subject: [PATCH 27/31] braking curves added

---
 etcs/brakingCurves.py | 425 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 425 insertions(+)
 create mode 100644 etcs/brakingCurves.py

diff --git a/etcs/brakingCurves.py b/etcs/brakingCurves.py
new file mode 100644
index 0000000..ca88af5
--- /dev/null
+++ b/etcs/brakingCurves.py
@@ -0,0 +1,425 @@
+from bisect import bisect_right
+
+import numpy as np
+import pandas as pd
+from matplotlib import pyplot as plt
+
+from mseetc.track import Track
+
+
+def compute_A_brake_safe(trainBrakingData):
+
+    braking = trainBrakingData["A_brake_emergency [m/s^2]"]
+
+    velocities = braking["velocity [m/s]"]
+    A_emergency_values = braking["value [m/s^2]"]
+
+    K_dry_rst = trainBrakingData["K_dry_rst [-]"]
+    M_NVAVADH = trainBrakingData["M_NVAVADH [-]"]
+    K_wet_rst = trainBrakingData["K_wet_rst [-]"]
+
+    K_wet_corr = K_wet_rst + M_NVAVADH * (1 - K_wet_rst)
+
+    A_brake_safe_values = [
+        A_emergency * K_dry_rst * K_wet_corr
+        for A_emergency in A_emergency_values
+    ]
+
+    return {
+            "velocity [m/s]": velocities,
+            "value [m/s^2]": A_brake_safe_values,
+    }
+
+
+def compute_A_gradient(currentPosition, gradients):
+
+    positions = gradients.index.values
+    gradients = gradients["Gradient [permil]"].values
+
+    idx = bisect_right(positions, currentPosition)
+    idx = max(0, min(idx, len(positions) - 1))
+    gradient = gradients[idx-1]
+
+    A_gradient = 9.81 * gradient * 0.001
+
+    return A_gradient
+
+
+def compute_braking_curve(braking_profile, gradients, target_position, permittedVelocity, dt=0.1):
+
+    max_velocity = permittedVelocity * 1.4
+
+    positions = [target_position]
+    velocities = [0.0]
+
+    threshold_velocities = list(braking_profile["velocity [m/s]"])
+    braking_values = list(braking_profile["value [m/s^2]"])
+
+    threshold_velocities.append(max_velocity)
+
+    for idx in range(len(threshold_velocities) - 1):
+
+        threshold_velocity = threshold_velocities[idx + 1]
+        A_brake = braking_values[idx]
+
+        while True:
+
+            A_gradient = compute_A_gradient(positions[-1], gradients)
+
+            v_old = velocities[-1]
+            v_new = v_old - (A_brake - A_gradient) * dt
+            x_new = positions[-1] - 0.5 * (v_new + v_old) * dt
+
+            positions.append(x_new)
+            velocities.append(v_new)
+
+            if v_new >= threshold_velocity or v_new >= max_velocity:
+
+                break
+
+        if velocities[-1] >= max_velocity:
+
+            break
+
+    curve = pd.DataFrame(
+        {"Velocity [m/s]": velocities[::-1]},
+        index=positions[::-1],
+    )
+
+    curve.index.name = "Position [m]"
+
+    return curve
+
+
+def compute_EBI_curve(EBD_curve, trainBrakingData):
+
+    positionsEBD = EBD_curve.index.to_numpy()
+    velocitiesEBD = EBD_curve["Velocity [m/s]"].to_numpy()
+
+    T_traction = trainBrakingData["T_traction [s]"]
+    T_be = trainBrakingData["T_be [s]"]
+    Kt_int = trainBrakingData["Kt_int [-]"]
+    v_uncertainty = trainBrakingData["v_uncertainty [%]"] * 0.01
+
+    t_be = T_be * Kt_int
+    T_berem = max(t_be - T_traction, 0)
+
+    positionsEBI = []
+    velocitiesEBI = []
+
+    for pos, vel in zip(positionsEBD, velocitiesEBD):
+
+        A_est1 = 0.1 # todo
+        A_est2 = min(A_est1, 0.4)
+
+        V_est = (vel - A_est1*T_traction - A_est2*T_berem)/(1+v_uncertainty)
+        V_est = max(V_est, 0.0)
+        velocitiesEBI.append(V_est)
+
+        V_delta_0 = V_est * v_uncertainty
+        V_delta1 = A_est1 * T_traction
+        V_delta2 = A_est2 * T_berem
+        D_bec = T_traction * (V_est+V_delta_0+0.5*V_delta1) + T_berem * (V_est+V_delta_0+V_delta1+0.5*V_delta2)
+        positionsEBI.append(pos - D_bec)
+
+    EBI_curve = pd.DataFrame(
+        {"Velocity [m/s]": velocitiesEBI},
+        index=positionsEBI,
+    )
+
+    EBI_curve.index.name = "Position [m]"
+
+    return EBI_curve
+
+
+def shift_curve_by_time(dfCurve, timeShift):
+
+    positionsOriginal = dfCurve.index.to_numpy()
+    velocitiesOriginal = dfCurve["Velocity [m/s]"].to_numpy()
+
+    velocitiesShifted = velocitiesOriginal.copy()
+    positionsShifted = positionsOriginal - velocitiesShifted * timeShift
+
+    dfCurveShifted = pd.DataFrame(
+        {"Velocity [m/s]": velocitiesShifted},
+        index=positionsShifted,
+    )
+
+    dfCurveShifted.index.name = "Position [m]"
+
+    return dfCurveShifted
+
+
+
+def compute_SBI_curve(SBI1_curve, SBI2_curve):
+
+    positionsSBI1 = SBI1_curve.index.to_numpy()
+    velocitiesSBI1 = SBI1_curve["Velocity [m/s]"].to_numpy()
+
+    positionsSBI2 = SBI2_curve.index.to_numpy()
+    velocitiesSBI2 = SBI2_curve["Velocity [m/s]"].to_numpy()
+
+    # Use only the overlapping position range
+    minPosition = max(positionsSBI1.min(), positionsSBI2.min())
+    maxPosition = min(positionsSBI1.max(), positionsSBI2.max())
+
+    step = 10.0  # [m]
+    positionsSBI = np.arange(minPosition, maxPosition, step)
+
+    velocitiesSBI1_interpol = np.interp(positionsSBI, positionsSBI1, velocitiesSBI1)
+    velocitiesSBI2_interpol = np.interp(positionsSBI, positionsSBI2, velocitiesSBI2)
+
+    velocitiesSBI = np.minimum(velocitiesSBI1_interpol, velocitiesSBI2_interpol)
+
+    SBI_curve = pd.DataFrame(
+        {"Velocity [m/s]": velocitiesSBI},
+        index=positionsSBI,
+    )
+
+    SBI_curve.index.name = "Position [m]"
+
+    return SBI_curve
+
+
+def getCurveStyles():
+
+    curveStyles = {
+        "EBD": {"color": "blue", "linestyle": "-", "linewidth": 2.0},
+        "EBI": {"color": "blue", "linestyle": ":", "linewidth": 1.5},
+        "SBD": {"color": "green", "linestyle": "-", "linewidth": 2.0},
+        "SBI1": {"color": "green", "linestyle": "--", "linewidth": 1.5},
+        "SBI2": {"color": "blue", "linestyle": "--", "linewidth": 1.5},
+        "SBI": {"color": "red", "linestyle": "-", "linewidth": 2.0},
+        "W": {"color": "orange", "linestyle": "-", "linewidth": 2.0},
+        "P": {"color": "grey", "linestyle": "-", "linewidth": 2.0},
+        "I": {"color": "gold", "linestyle": "-", "linewidth": 2.0},
+    }
+
+    return curveStyles
+
+
+def plotCurves(curves):
+
+    curveStyles = getCurveStyles()
+
+    fig, ax = plt.subplots(figsize=(16, 8))
+
+    for name, curve in curves.items():
+
+        style = {}
+
+        if curveStyles is not None and name in curveStyles:
+            style = curveStyles[name]
+
+        ax.plot(curve.index.values / 1000, curve["Velocity [m/s]"] * 3.6, label=name, **style)
+
+    ax.set_title("ETCS Braking Curves")
+    ax.set_xlabel("Position [km]")
+    ax.set_ylabel("Velocity [km/h]")
+    ax.grid(True, which="both", linestyle="--", alpha=0.5)
+    ax.legend(loc="upper right")
+    ax.figure.tight_layout()
+
+    plt.show()
+
+
+def computeCeilingSpeedLimits(V_permitted_mps):
+
+    dV_ebi_min = 7.5/3.6
+    dV_ebi_max = 15.0/3.6
+    V_ebi_min = 110.0/3.6
+    V_ebi_max = 210.0/3.6
+
+    C_ebi = (dV_ebi_max - dV_ebi_min) / (V_ebi_max - V_ebi_min)
+
+    if V_permitted_mps <= V_ebi_min:
+
+        dV_ebi = dV_ebi_min
+
+    else:
+        dV_ebi = min(dV_ebi_min + C_ebi * (V_permitted_mps - V_ebi_min),dV_ebi_max,)
+
+    dV_warning = 0.5 * dV_ebi
+    dV_sbi = 0.75 * dV_ebi
+
+    return {
+        "Warning [m/s]": V_permitted_mps + dV_warning,
+        "SBI [m/s]": V_permitted_mps + dV_sbi,
+        "EBI [m/s]": V_permitted_mps + dV_ebi,
+    }
+
+
+def addStartPointToCurve(curve, velocity, start_position):
+
+    start_point = pd.DataFrame(
+        {"Velocity [m/s]": [velocity,velocity]},
+        index=[start_position, curve.index.to_numpy(dtype=float)[0]-10],
+    )
+    start_point.index.name = "Position [m]"
+
+    return pd.concat([start_point, curve])
+
+
+def trimCurveToMaxVelocity(curve, maxVelocity):
+
+    velocities = curve["Velocity [m/s]"].to_numpy(dtype=float)
+    keep_mask = velocities <= maxVelocity
+
+    return curve[keep_mask].copy()
+
+
+def trimCurveFromMinPosition(curve, minPosition):
+
+    positions = curve.index.to_numpy(dtype=float)
+    keep_mask = positions >= minPosition
+
+    return curve[keep_mask].copy()
+
+
+def postProcessCurves(curves, permittedSpeed):
+
+    start_position = curves["P"].index.to_numpy(dtype=float)[-1] - 3000
+
+    speedLimits = computeCeilingSpeedLimits(permittedSpeed)
+
+    curves["EBI"] = trimCurveToMaxVelocity(curves["EBI"], speedLimits["EBI [m/s]"])
+    curves["EBI"] = addStartPointToCurve(curves["EBI"], speedLimits["EBI [m/s]"], start_position)
+
+    curves["SBI"] = trimCurveToMaxVelocity(curves["SBI"], speedLimits["SBI [m/s]"])
+    curves["SBI"] = addStartPointToCurve(curves["SBI"], speedLimits["SBI [m/s]"], start_position)
+
+    curves["W"] = trimCurveToMaxVelocity(curves["W"], speedLimits["Warning [m/s]"])
+    curves["W"] = addStartPointToCurve(curves["W"], speedLimits["Warning [m/s]"], start_position)
+
+    curves["P"] = trimCurveToMaxVelocity(curves["P"], permittedSpeed)
+    curves["P"] = addStartPointToCurve(curves["P"], permittedSpeed, start_position)
+
+    curves["I"] = trimCurveToMaxVelocity(curves["I"], permittedSpeed)
+
+
+    curves["EBD"] = trimCurveFromMinPosition(curves["EBD"], curves["EBI"].index.to_numpy(dtype=float)[2])
+
+    curves["SBI2"] = trimCurveFromMinPosition(curves["SBI2"], curves["EBI"].index.to_numpy(dtype=float)[2])
+
+    curves["SBD"] = trimCurveFromMinPosition(curves["SBD"], curves["SBI"].index.to_numpy(dtype=float)[2])
+
+    curves["SBI1"] = trimCurveFromMinPosition(curves["SBI1"], curves["SBI"].index.to_numpy(dtype=float)[2])
+
+    return curves
+
+
+def computeStepApproximation(curve, step=50):
+
+    positions = curve.index.to_numpy(dtype=float)
+    velocities = curve["Velocity [m/s]"].to_numpy(dtype=float)
+
+    samplePositions = np.arange(positions[1], positions[-1], step)
+    sampleVelocities = np.interp(samplePositions, positions, velocities)
+
+    approximation = pd.DataFrame(
+        {"Velocity [m/s]": sampleVelocities},
+        index=samplePositions,
+    )
+
+    approximation.index.name = "Position [m]"
+
+    return approximation
+
+
+def plotApproximation(approximation, curve, name):
+
+    curveStyles = getCurveStyles()
+
+    fig, ax = plt.subplots(figsize=(16, 8))
+
+    style = {}
+
+    if curveStyles is not None and name in curveStyles:
+        style = curveStyles[name]
+
+    ax.plot(curve.index.values / 1000, curve["Velocity [m/s]"] * 3.6, label=name, **style)
+    ax.step(approximation.index.values / 1000, approximation["Velocity [m/s]"] * 3.6, label="Step Approxmiation")
+
+    ax.set_title("ETCS Braking Curves")
+    ax.set_xlabel("Position [km]")
+    ax.set_ylabel("Velocity [km/h]")
+    ax.grid(True, which="both", linestyle="--", alpha=0.5)
+    ax.legend(loc="upper right")
+    ax.figure.tight_layout()
+
+    plt.show()
+
+
+if __name__ == '__main__':
+
+    # ETCS Constants
+    T_warning = 2  # [s]
+    T_driver = 4  # [s]
+
+    # Convention:
+    # Braking accelerations are stored as negative values.
+    # Gradient acceleration is positive for uphill and negative for downhill.
+
+    trainBrakingData = {
+        "A_brake_emergency [m/s^2]": {
+            "velocity [m/s]": [0, 20, 40, 60],
+            "value [m/s^2]": [-0.9, -0.85, -0.8, -0.75],
+        },
+        "A_brake_service [m/s^2]": {
+            "velocity [m/s]": [0, 20, 40, 60],
+            "value [m/s^2]": [-0.5, -0.45, -0.4, -0.35],
+        },
+        "K_dry_rst [-]": 0.8,
+        "M_NVAVADH [-]": 0,
+        "K_wet_rst [-]": 0.9,
+        "T_traction [s]": 1,
+        "T_be [s]": 4,
+        "Kt_int [-]": 1.15,
+        "v_uncertainty [%]": 2.98,
+        "T_bs [s]": 3,
+        "T_bs1 [s]": 3,
+        "T_bs2 [s]": 3,
+    }
+
+    track = Track(config={'id': 'CH_StGallen_Wil'}, pathJSON='../tracks')
+
+    targetPosition = 5000  # [m]
+    overlap = 100  # [m]
+    permittedSpeed = 160 / 3.6  # [m/s]
+    targetSpeed = 0  # [m/s]
+
+    SvL = targetPosition + overlap
+    EoA = targetPosition
+
+    assert targetPosition > 0 and targetPosition < track.length
+
+
+    df_A_brake_safe = compute_A_brake_safe(trainBrakingData)
+
+    curves = {}
+
+    curves["EBD"] = compute_braking_curve(df_A_brake_safe, track.gradients, SvL, permittedSpeed)
+
+    curves["EBI"] = compute_EBI_curve(curves["EBD"], trainBrakingData)
+
+    curves["SBI2"] = shift_curve_by_time(curves["EBI"], trainBrakingData["T_bs2 [s]"])
+
+    curves["SBD"] = compute_braking_curve(trainBrakingData["A_brake_service [m/s^2]"], track.gradients, EoA, permittedSpeed)
+
+    curves["SBI1"] = shift_curve_by_time(curves["SBD"], trainBrakingData["T_bs1 [s]"])
+
+    curves["SBI"] = compute_SBI_curve(curves["SBI1"], curves["SBI2"])
+
+    curves["W"] = shift_curve_by_time(curves["SBI"], T_warning)
+
+    curves["P"] = shift_curve_by_time(curves["SBI"], T_driver)
+
+    T_indication = max(0.8 * trainBrakingData["T_bs [s]"], 5) + T_driver
+    curves["I"] = shift_curve_by_time(curves["P"], T_indication)
+
+    curves = postProcessCurves(curves, permittedSpeed)
+
+    plotCurves(curves)
+
+    approximation = computeStepApproximation(curves["P"])
+    plotApproximation(approximation, curves["P"], "P")

From 390c11653d7201a42a02823482c0172b10ff0b7b Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Fri, 5 Jun 2026 15:15:54 +0200
Subject: [PATCH 28/31] etcs file added

---
 etcs/brakingCurves.py => mseetc/etcs.py | 132 ++++++++++++++++++++----
 1 file changed, 111 insertions(+), 21 deletions(-)
 rename etcs/brakingCurves.py => mseetc/etcs.py (74%)

diff --git a/etcs/brakingCurves.py b/mseetc/etcs.py
similarity index 74%
rename from etcs/brakingCurves.py
rename to mseetc/etcs.py
index ca88af5..d30df0d 100644
--- a/etcs/brakingCurves.py
+++ b/mseetc/etcs.py
@@ -1,4 +1,6 @@
 from bisect import bisect_right
+from idlelib.editor import prepstr
+from multiprocessing.spawn import prepare
 
 import numpy as np
 import pandas as pd
@@ -38,27 +40,34 @@ def compute_A_gradient(currentPosition, gradients):
 
     idx = bisect_right(positions, currentPosition)
     idx = max(0, min(idx, len(positions) - 1))
-    gradient = gradients[idx-1]
 
-    A_gradient = 9.81 * gradient * 0.001
+    if idx == 0:
 
+        return 0
+
+    gradient = gradients[idx-1]
+    A_gradient = 9.81 * gradient * 0.001
     return A_gradient
 
 
-def compute_braking_curve(braking_profile, gradients, target_position, permittedVelocity, dt=0.1):
+def compute_braking_curve(braking_profile, gradients, target_position, permittedVelocity, targetVerlocity, dt=0.1):
 
     max_velocity = permittedVelocity * 1.4
 
     positions = [target_position]
-    velocities = [0.0]
+    velocities = [targetVerlocity]
 
     threshold_velocities = list(braking_profile["velocity [m/s]"])
     braking_values = list(braking_profile["value [m/s^2]"])
 
-    threshold_velocities.append(max_velocity)
+    threshold_velocities.append(200)  # basically inf velocity
 
     for idx in range(len(threshold_velocities) - 1):
 
+        if targetVerlocity > threshold_velocities[idx + 1]:
+
+            continue
+
         threshold_velocity = threshold_velocities[idx + 1]
         A_brake = braking_values[idx]
 
@@ -91,7 +100,7 @@ def compute_braking_curve(braking_profile, gradients, target_position, permitted
     return curve
 
 
-def compute_EBI_curve(EBD_curve, trainBrakingData):
+def compute_EBI_curve(EBD_curve, trainBrakingData, targetSpeed):
 
     positionsEBD = EBD_curve.index.to_numpy()
     velocitiesEBD = EBD_curve["Velocity [m/s]"].to_numpy()
@@ -114,6 +123,14 @@ def compute_EBI_curve(EBD_curve, trainBrakingData):
 
         V_est = (vel - A_est1*T_traction - A_est2*T_berem)/(1+v_uncertainty)
         V_est = max(V_est, 0.0)
+
+        if V_est < targetSpeed:
+
+            velocitiesEBI.append(targetSpeed)
+            positionsEBI.append(positionsEBI[-1]+1)
+
+            break
+
         velocitiesEBI.append(V_est)
 
         V_delta_0 = V_est * v_uncertainty
@@ -164,7 +181,8 @@ def compute_SBI_curve(SBI1_curve, SBI2_curve):
     maxPosition = min(positionsSBI1.max(), positionsSBI2.max())
 
     step = 10.0  # [m]
-    positionsSBI = np.arange(minPosition, maxPosition, step)
+    positionsSBI = np.arange(minPosition, maxPosition + step, step)
+    positionsSBI[-1] = maxPosition
 
     velocitiesSBI1_interpol = np.interp(positionsSBI, positionsSBI1, velocitiesSBI1)
     velocitiesSBI2_interpol = np.interp(positionsSBI, positionsSBI2, velocitiesSBI2)
@@ -198,12 +216,18 @@ def getCurveStyles():
     return curveStyles
 
 
-def plotCurves(curves):
+def plotCurves(curves, targetPosition, permittedSpeed, targetSpeed, prePottingDistance, postPlottingDistance):
 
     curveStyles = getCurveStyles()
 
     fig, ax = plt.subplots(figsize=(16, 8))
 
+    ax.step(
+        np.array([targetPosition - prePottingDistance, targetPosition, targetPosition + postPlottingDistance]) / 1000,
+        np.array([permittedSpeed, permittedSpeed, targetSpeed]) * 3.6,
+        label="Speed limit", color="black", linewidth= 2.0
+    )
+
     for name, curve in curves.items():
 
         style = {}
@@ -251,15 +275,30 @@ def computeCeilingSpeedLimits(V_permitted_mps):
 
 def addStartPointToCurve(curve, velocity, start_position):
 
+    delta_s = curve.index.to_numpy(dtype=float)[1] - curve.index.to_numpy(dtype=float)[0]
+
     start_point = pd.DataFrame(
-        {"Velocity [m/s]": [velocity,velocity]},
-        index=[start_position, curve.index.to_numpy(dtype=float)[0]-10],
+        {"Velocity [m/s]": [velocity, velocity]},
+        index=[start_position, curve.index.to_numpy(dtype=float)[0] - delta_s],
     )
     start_point.index.name = "Position [m]"
 
     return pd.concat([start_point, curve])
 
 
+def addEndPointToCurve(curve, velocity, end_position):
+
+    delta_s = curve.index.to_numpy(dtype=float)[-1] - curve.index.to_numpy(dtype=float)[-2]
+
+    end_point = pd.DataFrame(
+        {"Velocity [m/s]": [velocity, velocity]},
+        index=[curve.index.to_numpy(dtype=float)[-1] + delta_s, end_position],
+    )
+    end_point.index.name = "Position [m]"
+
+    return pd.concat([curve, end_point])
+
+
 def trimCurveToMaxVelocity(curve, maxVelocity):
 
     velocities = curve["Velocity [m/s]"].to_numpy(dtype=float)
@@ -268,6 +307,14 @@ def trimCurveToMaxVelocity(curve, maxVelocity):
     return curve[keep_mask].copy()
 
 
+def trimCurveFromMinVelocity(curve, minVelocity):
+
+    velocities = curve["Velocity [m/s]"].to_numpy(dtype=float)
+    keep_mask = velocities >= minVelocity
+
+    return curve[keep_mask].copy()
+
+
 def trimCurveFromMinPosition(curve, minPosition):
 
     positions = curve.index.to_numpy(dtype=float)
@@ -276,9 +323,9 @@ def trimCurveFromMinPosition(curve, minPosition):
     return curve[keep_mask].copy()
 
 
-def postProcessCurves(curves, permittedSpeed):
+def postPreProcessCurves(curves, targetPosition, permittedSpeed, plottingDistance):
 
-    start_position = curves["P"].index.to_numpy(dtype=float)[-1] - 3000
+    start_position = targetPosition - plottingDistance
 
     speedLimits = computeCeilingSpeedLimits(permittedSpeed)
 
@@ -308,6 +355,37 @@ def postProcessCurves(curves, permittedSpeed):
     return curves
 
 
+def postPostProcessCurves(curves, targetPosition, targetSpeed, postPlottingDistance):
+
+    end_position = targetPosition + postPlottingDistance
+
+    speedLimits = computeCeilingSpeedLimits(targetSpeed)
+
+    curves["EBI"] = trimCurveFromMinVelocity(curves["EBI"], speedLimits["EBI [m/s]"])
+    curves["EBI"] = addEndPointToCurve(curves["EBI"], speedLimits["EBI [m/s]"], end_position)
+
+    curves["SBI"] = trimCurveFromMinVelocity(curves["SBI"], speedLimits["SBI [m/s]"])
+    curves["SBI"] = addEndPointToCurve(curves["SBI"], speedLimits["SBI [m/s]"], end_position)
+
+    curves["W"] = trimCurveFromMinVelocity(curves["W"], speedLimits["Warning [m/s]"])
+    curves["W"] = addEndPointToCurve(curves["W"], speedLimits["Warning [m/s]"], end_position)
+
+    curves["P"] = addEndPointToCurve(curves["P"], targetSpeed, end_position)
+
+    curves["I"] = addEndPointToCurve(curves["I"], targetSpeed, curves["P"].index.to_numpy(dtype=float)[-4])
+
+
+    curves["EBD"] = trimCurveFromMinVelocity(curves["EBD"], speedLimits["EBI [m/s]"])
+
+    curves["SBI2"] = trimCurveFromMinVelocity(curves["SBI2"], speedLimits["SBI [m/s]"])
+
+    curves["SBD"] = trimCurveFromMinVelocity(curves["SBD"], speedLimits["EBI [m/s]"])
+
+    curves["SBI1"] = trimCurveFromMinVelocity(curves["SBI1"], speedLimits["SBI [m/s]"])
+
+    return curves
+
+
 def computeStepApproximation(curve, step=50):
 
     positions = curve.index.to_numpy(dtype=float)
@@ -386,25 +464,34 @@ def plotApproximation(approximation, curve, name):
     targetPosition = 5000  # [m]
     overlap = 100  # [m]
     permittedSpeed = 160 / 3.6  # [m/s]
-    targetSpeed = 0  # [m/s]
+    targetSpeed = 24  # [m/s]
+
+    prepPlottingDistance = 3000  # [m]
+    postPlottingDistance = 1000  # [m]
 
     SvL = targetPosition + overlap
     EoA = targetPosition
 
+    assert permittedSpeed >= 0 and permittedSpeed < 400 / 3.6
+
     assert targetPosition > 0 and targetPosition < track.length
 
+    assert targetSpeed >= 0 and targetSpeed < permittedSpeed
+
 
     df_A_brake_safe = compute_A_brake_safe(trainBrakingData)
 
+    T_indication = max(0.8 * trainBrakingData["T_bs [s]"], 5) + T_driver
+
     curves = {}
 
-    curves["EBD"] = compute_braking_curve(df_A_brake_safe, track.gradients, SvL, permittedSpeed)
+    curves["EBD"] = compute_braking_curve(df_A_brake_safe, track.gradients, SvL, permittedSpeed, targetSpeed)
 
-    curves["EBI"] = compute_EBI_curve(curves["EBD"], trainBrakingData)
+    curves["EBI"] = compute_EBI_curve(curves["EBD"], trainBrakingData, targetSpeed)
 
     curves["SBI2"] = shift_curve_by_time(curves["EBI"], trainBrakingData["T_bs2 [s]"])
 
-    curves["SBD"] = compute_braking_curve(trainBrakingData["A_brake_service [m/s^2]"], track.gradients, EoA, permittedSpeed)
+    curves["SBD"] = compute_braking_curve(trainBrakingData["A_brake_service [m/s^2]"], track.gradients, EoA, permittedSpeed, targetSpeed)
 
     curves["SBI1"] = shift_curve_by_time(curves["SBD"], trainBrakingData["T_bs1 [s]"])
 
@@ -414,12 +501,15 @@ def plotApproximation(approximation, curve, name):
 
     curves["P"] = shift_curve_by_time(curves["SBI"], T_driver)
 
-    T_indication = max(0.8 * trainBrakingData["T_bs [s]"], 5) + T_driver
     curves["I"] = shift_curve_by_time(curves["P"], T_indication)
 
-    curves = postProcessCurves(curves, permittedSpeed)
+    curves = postPreProcessCurves(curves, targetPosition, permittedSpeed, prepPlottingDistance)
+
+    if targetSpeed > 0:
+
+        curves = postPostProcessCurves(curves, targetPosition, targetSpeed, postPlottingDistance)
 
-    plotCurves(curves)
+    plotCurves(curves, targetPosition, permittedSpeed, targetSpeed, prepPlottingDistance, postPlottingDistance)
 
-    approximation = computeStepApproximation(curves["P"])
-    plotApproximation(approximation, curves["P"], "P")
+    # approximation = computeStepApproximation(curves["P"])
+    # plotApproximation(approximation, curves["P"], "P")

From 52b83c7276078f064e658e80984ed8d5625404a6 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Fri, 5 Jun 2026 17:02:37 +0200
Subject: [PATCH 29/31] refactoring etcs

---
 mseetc/etcs.py | 649 ++++++++++++++++++++++++++-----------------------
 1 file changed, 343 insertions(+), 306 deletions(-)

diff --git a/mseetc/etcs.py b/mseetc/etcs.py
index d30df0d..1b1e457 100644
--- a/mseetc/etcs.py
+++ b/mseetc/etcs.py
@@ -1,6 +1,5 @@
 from bisect import bisect_right
-from idlelib.editor import prepstr
-from multiprocessing.spawn import prepare
+from dataclasses import dataclass
 
 import numpy as np
 import pandas as pd
@@ -9,430 +8,506 @@
 from mseetc.track import Track
 
 
-def compute_A_brake_safe(trainBrakingData):
+def shiftCurveByTime(dfCurve, timeShift):
 
-    braking = trainBrakingData["A_brake_emergency [m/s^2]"]
+    positionsOriginal = dfCurve.index.to_numpy()
+    velocitiesOriginal = dfCurve["Velocity [m/s]"].to_numpy()
+
+    velocitiesShifted = velocitiesOriginal.copy()
+    positionsShifted = positionsOriginal - velocitiesShifted * timeShift
+
+    dfCurveShifted = pd.DataFrame(
+        {"Velocity [m/s]": velocitiesShifted},
+        index=positionsShifted,
+    )
 
-    velocities = braking["velocity [m/s]"]
-    A_emergency_values = braking["value [m/s^2]"]
+    dfCurveShifted.index.name = "Position [m]"
+
+    return dfCurveShifted
 
-    K_dry_rst = trainBrakingData["K_dry_rst [-]"]
-    M_NVAVADH = trainBrakingData["M_NVAVADH [-]"]
-    K_wet_rst = trainBrakingData["K_wet_rst [-]"]
 
-    K_wet_corr = K_wet_rst + M_NVAVADH * (1 - K_wet_rst)
+def computeCeilingSpeedLimits(V_permitted_mps):
+
+    dV_ebi_min = 7.5/3.6
+    dV_ebi_max = 15.0/3.6
+    V_ebi_min = 110.0/3.6
+    V_ebi_max = 210.0/3.6
 
-    A_brake_safe_values = [
-        A_emergency * K_dry_rst * K_wet_corr
-        for A_emergency in A_emergency_values
-    ]
+    C_ebi = (dV_ebi_max - dV_ebi_min) / (V_ebi_max - V_ebi_min)
+
+    if V_permitted_mps <= V_ebi_min:
+
+        dV_ebi = dV_ebi_min
+
+    else:
+        dV_ebi = min(dV_ebi_min + C_ebi * (V_permitted_mps - V_ebi_min),dV_ebi_max,)
+
+    dV_warning = 0.5 * dV_ebi
+    dV_sbi = 0.75 * dV_ebi
 
     return {
-            "velocity [m/s]": velocities,
-            "value [m/s^2]": A_brake_safe_values,
+        "Warning [m/s]": V_permitted_mps + dV_warning,
+        "SBI [m/s]": V_permitted_mps + dV_sbi,
+        "EBI [m/s]": V_permitted_mps + dV_ebi,
     }
 
 
-def compute_A_gradient(currentPosition, gradients):
+def addStartPointToCurve(curve, velocity, start_position):
 
-    positions = gradients.index.values
-    gradients = gradients["Gradient [permil]"].values
+    delta_s = curve.index.to_numpy(dtype=float)[1] - curve.index.to_numpy(dtype=float)[0]
 
-    idx = bisect_right(positions, currentPosition)
-    idx = max(0, min(idx, len(positions) - 1))
+    start_point = pd.DataFrame(
+        {"Velocity [m/s]": [velocity, velocity]},
+        index=[start_position, curve.index.to_numpy(dtype=float)[0] - delta_s],
+    )
+    start_point.index.name = "Position [m]"
 
-    if idx == 0:
+    return pd.concat([start_point, curve])
 
-        return 0
 
-    gradient = gradients[idx-1]
-    A_gradient = 9.81 * gradient * 0.001
-    return A_gradient
+def addEndPointToCurve(curve, velocity, end_position):
 
+    delta_s = curve.index.to_numpy(dtype=float)[-1] - curve.index.to_numpy(dtype=float)[-2]
 
-def compute_braking_curve(braking_profile, gradients, target_position, permittedVelocity, targetVerlocity, dt=0.1):
+    end_point = pd.DataFrame(
+        {"Velocity [m/s]": [velocity, velocity]},
+        index=[curve.index.to_numpy(dtype=float)[-1] + delta_s, end_position],
+    )
+    end_point.index.name = "Position [m]"
 
-    max_velocity = permittedVelocity * 1.4
+    return pd.concat([curve, end_point])
 
-    positions = [target_position]
-    velocities = [targetVerlocity]
 
-    threshold_velocities = list(braking_profile["velocity [m/s]"])
-    braking_values = list(braking_profile["value [m/s^2]"])
+def trimCurveToMaxVelocity(curve, maxVelocity):
 
-    threshold_velocities.append(200)  # basically inf velocity
+    velocities = curve["Velocity [m/s]"].to_numpy(dtype=float)
+    keep_mask = velocities <= maxVelocity
 
-    for idx in range(len(threshold_velocities) - 1):
+    return curve[keep_mask].copy()
 
-        if targetVerlocity > threshold_velocities[idx + 1]:
 
-            continue
+def trimCurveFromMinVelocity(curve, minVelocity):
 
-        threshold_velocity = threshold_velocities[idx + 1]
-        A_brake = braking_values[idx]
+    velocities = curve["Velocity [m/s]"].to_numpy(dtype=float)
+    keep_mask = velocities >= minVelocity
 
-        while True:
+    return curve[keep_mask].copy()
 
-            A_gradient = compute_A_gradient(positions[-1], gradients)
 
-            v_old = velocities[-1]
-            v_new = v_old - (A_brake - A_gradient) * dt
-            x_new = positions[-1] - 0.5 * (v_new + v_old) * dt
+def trimCurveFromMinPosition(curve, minPosition):
 
-            positions.append(x_new)
-            velocities.append(v_new)
+    positions = curve.index.to_numpy(dtype=float)
+    keep_mask = positions >= minPosition
 
-            if v_new >= threshold_velocity or v_new >= max_velocity:
+    return curve[keep_mask].copy()
 
-                break
 
-        if velocities[-1] >= max_velocity:
+@dataclass(frozen=True)
+class BrakingTarget:
+    position: float  # [m]
+    overlap: float  # [m]
+    permittedVelocity: float  # [m/s]
+    targetVelocity: float  # [m/s]
 
-            break
+    @property
+    def EoA(self):
+        return self.position
 
-    curve = pd.DataFrame(
-        {"Velocity [m/s]": velocities[::-1]},
-        index=positions[::-1],
-    )
+    @property
+    def SvL(self):
+        return self.position + self.overlap
 
-    curve.index.name = "Position [m]"
 
-    return curve
+class EtcsBrakingCurveCalculator:
 
+    def __init__(self, trainBrakingData, track, distancePre=3000, distancePost=1000):
 
-def compute_EBI_curve(EBD_curve, trainBrakingData, targetSpeed):
+        self.trainBrakingData = trainBrakingData
+        self.track = track
 
-    positionsEBD = EBD_curve.index.to_numpy()
-    velocitiesEBD = EBD_curve["Velocity [m/s]"].to_numpy()
+        self.dt = 0.1  # [s]
+        self.distancePre = distancePre
+        self.distancePost = distancePost
 
-    T_traction = trainBrakingData["T_traction [s]"]
-    T_be = trainBrakingData["T_be [s]"]
-    Kt_int = trainBrakingData["Kt_int [-]"]
-    v_uncertainty = trainBrakingData["v_uncertainty [%]"] * 0.01
+        # ETCS Constants
+        self.T_warning = 2.0  # [s]
+        self.T_driver = 4.0  # [s]
 
-    t_be = T_be * Kt_int
-    T_berem = max(t_be - T_traction, 0)
+        self.curveStyles = {
+            "EBD": {"color": "blue", "linestyle": "-", "linewidth": 2.0},
+            "EBI": {"color": "blue", "linestyle": ":", "linewidth": 1.5},
+            "SBD": {"color": "green", "linestyle": "-", "linewidth": 2.0},
+            "SBI1": {"color": "green", "linestyle": "--", "linewidth": 1.5},
+            "SBI2": {"color": "blue", "linestyle": "--", "linewidth": 1.5},
+            "SBI": {"color": "red", "linestyle": "-", "linewidth": 2.0},
+            "W": {"color": "orange", "linestyle": "-", "linewidth": 2.0},
+            "P": {"color": "grey", "linestyle": "-", "linewidth": 2.0},
+            "I": {"color": "gold", "linestyle": "-", "linewidth": 2.0},
+        }
 
-    positionsEBI = []
-    velocitiesEBI = []
 
-    for pos, vel in zip(positionsEBD, velocitiesEBD):
+    def validateInput(self, target):
 
-        A_est1 = 0.1 # todo
-        A_est2 = min(A_est1, 0.4)
+        if not 0 <= target.permittedVelocity < 400 / 3.6:
+            raise ValueError("permittedVelocity must be between 0 and 400 km/h.")
 
-        V_est = (vel - A_est1*T_traction - A_est2*T_berem)/(1+v_uncertainty)
-        V_est = max(V_est, 0.0)
+        if not 0 < target.EoA < self.track.length:
+            raise ValueError("EoA must lie within the track length.")
 
-        if V_est < targetSpeed:
+        if not 0 < target.SvL < self.track.length:
+            raise ValueError("SvL must lie within the track length.")
 
-            velocitiesEBI.append(targetSpeed)
-            positionsEBI.append(positionsEBI[-1]+1)
+        if not 0 <= target.targetVelocity < target.permittedVelocity:
+            raise ValueError("targetVelocity must be lower than permittedVelocity.")
 
-            break
 
-        velocitiesEBI.append(V_est)
+    def compute_A_brake_safe(self):
 
-        V_delta_0 = V_est * v_uncertainty
-        V_delta1 = A_est1 * T_traction
-        V_delta2 = A_est2 * T_berem
-        D_bec = T_traction * (V_est+V_delta_0+0.5*V_delta1) + T_berem * (V_est+V_delta_0+V_delta1+0.5*V_delta2)
-        positionsEBI.append(pos - D_bec)
+        trainBrakingData = self.trainBrakingData
 
-    EBI_curve = pd.DataFrame(
-        {"Velocity [m/s]": velocitiesEBI},
-        index=positionsEBI,
-    )
+        braking = trainBrakingData["A_brake_emergency [m/s^2]"]
 
-    EBI_curve.index.name = "Position [m]"
+        velocities = braking["velocity [m/s]"]
+        A_emergency_values = braking["value [m/s^2]"]
 
-    return EBI_curve
+        K_dry_rst = trainBrakingData["K_dry_rst [-]"]
+        M_NVAVADH = trainBrakingData["M_NVAVADH [-]"]
+        K_wet_rst = trainBrakingData["K_wet_rst [-]"]
 
+        K_wet_corr = K_wet_rst + M_NVAVADH * (1 - K_wet_rst)
 
-def shift_curve_by_time(dfCurve, timeShift):
+        A_brake_safe_values = [
+            A_emergency * K_dry_rst * K_wet_corr
+            for A_emergency in A_emergency_values
+        ]
 
-    positionsOriginal = dfCurve.index.to_numpy()
-    velocitiesOriginal = dfCurve["Velocity [m/s]"].to_numpy()
+        return {
+            "velocity [m/s]": velocities,
+            "value [m/s^2]": A_brake_safe_values,
+        }
 
-    velocitiesShifted = velocitiesOriginal.copy()
-    positionsShifted = positionsOriginal - velocitiesShifted * timeShift
 
-    dfCurveShifted = pd.DataFrame(
-        {"Velocity [m/s]": velocitiesShifted},
-        index=positionsShifted,
-    )
+    def compute_A_gradient(self, currentPosition):
 
-    dfCurveShifted.index.name = "Position [m]"
+        positions = self.track.gradients.index.values
+        gradients = self.track.gradients["Gradient [permil]"].values
 
-    return dfCurveShifted
+        idx = bisect_right(positions, currentPosition) - 1
+        idx = max(0, min(idx, len(positions) - 1))
 
+        if idx == 0:
 
+            return 0  # [start of track has been exceeded]
 
-def compute_SBI_curve(SBI1_curve, SBI2_curve):
+        gradient = gradients[idx]
+        return 9.81 * gradient * 0.001
 
-    positionsSBI1 = SBI1_curve.index.to_numpy()
-    velocitiesSBI1 = SBI1_curve["Velocity [m/s]"].to_numpy()
 
-    positionsSBI2 = SBI2_curve.index.to_numpy()
-    velocitiesSBI2 = SBI2_curve["Velocity [m/s]"].to_numpy()
+    def computeBrakingCurve(self, brakingProfile, target_position, permittedVelocity, targetVelocity):
 
-    # Use only the overlapping position range
-    minPosition = max(positionsSBI1.min(), positionsSBI2.min())
-    maxPosition = min(positionsSBI1.max(), positionsSBI2.max())
+        max_velocity = permittedVelocity * 1.4
 
-    step = 10.0  # [m]
-    positionsSBI = np.arange(minPosition, maxPosition + step, step)
-    positionsSBI[-1] = maxPosition
+        positions = [target_position]
+        velocities = [targetVelocity]
 
-    velocitiesSBI1_interpol = np.interp(positionsSBI, positionsSBI1, velocitiesSBI1)
-    velocitiesSBI2_interpol = np.interp(positionsSBI, positionsSBI2, velocitiesSBI2)
+        threshold_velocities = list(brakingProfile["velocity [m/s]"])
+        braking_values = list(brakingProfile["value [m/s^2]"])
 
-    velocitiesSBI = np.minimum(velocitiesSBI1_interpol, velocitiesSBI2_interpol)
+        threshold_velocities.append(200)  # basically inf velocity
 
-    SBI_curve = pd.DataFrame(
-        {"Velocity [m/s]": velocitiesSBI},
-        index=positionsSBI,
-    )
+        for idx in range(len(threshold_velocities) - 1):
 
-    SBI_curve.index.name = "Position [m]"
+            if targetVelocity > threshold_velocities[idx + 1]:
+                continue
 
-    return SBI_curve
+            threshold_velocity = threshold_velocities[idx + 1]
+            A_brake = braking_values[idx]
 
+            while True:
 
-def getCurveStyles():
+                A_gradient = self.compute_A_gradient(positions[-1])
 
-    curveStyles = {
-        "EBD": {"color": "blue", "linestyle": "-", "linewidth": 2.0},
-        "EBI": {"color": "blue", "linestyle": ":", "linewidth": 1.5},
-        "SBD": {"color": "green", "linestyle": "-", "linewidth": 2.0},
-        "SBI1": {"color": "green", "linestyle": "--", "linewidth": 1.5},
-        "SBI2": {"color": "blue", "linestyle": "--", "linewidth": 1.5},
-        "SBI": {"color": "red", "linestyle": "-", "linewidth": 2.0},
-        "W": {"color": "orange", "linestyle": "-", "linewidth": 2.0},
-        "P": {"color": "grey", "linestyle": "-", "linewidth": 2.0},
-        "I": {"color": "gold", "linestyle": "-", "linewidth": 2.0},
-    }
+                v_old = velocities[-1]
+                v_new = v_old - (A_brake - A_gradient) * self.dt
+                x_new = positions[-1] - 0.5 * (v_new + v_old) * self.dt
 
-    return curveStyles
+                positions.append(x_new)
+                velocities.append(v_new)
 
+                if v_new >= threshold_velocity or v_new >= max_velocity:
+                    break
 
-def plotCurves(curves, targetPosition, permittedSpeed, targetSpeed, prePottingDistance, postPlottingDistance):
+            if velocities[-1] >= max_velocity:
+                break
 
-    curveStyles = getCurveStyles()
+        curve = pd.DataFrame(
+            {"Velocity [m/s]": velocities[::-1]},
+            index=positions[::-1],
+        )
 
-    fig, ax = plt.subplots(figsize=(16, 8))
+        curve.index.name = "Position [m]"
 
-    ax.step(
-        np.array([targetPosition - prePottingDistance, targetPosition, targetPosition + postPlottingDistance]) / 1000,
-        np.array([permittedSpeed, permittedSpeed, targetSpeed]) * 3.6,
-        label="Speed limit", color="black", linewidth= 2.0
-    )
+        return curve
 
-    for name, curve in curves.items():
 
-        style = {}
+    def compute_EBI_curve(self, EBD_curve, targetVelocity):
 
-        if curveStyles is not None and name in curveStyles:
-            style = curveStyles[name]
+        trainBrakingData = self.trainBrakingData
 
-        ax.plot(curve.index.values / 1000, curve["Velocity [m/s]"] * 3.6, label=name, **style)
+        T_traction = trainBrakingData["T_traction [s]"]
+        T_be = trainBrakingData["T_be [s]"]
+        Kt_int = trainBrakingData["Kt_int [-]"]
+        v_uncertainty = trainBrakingData["v_uncertainty [%]"] * 0.01
 
-    ax.set_title("ETCS Braking Curves")
-    ax.set_xlabel("Position [km]")
-    ax.set_ylabel("Velocity [km/h]")
-    ax.grid(True, which="both", linestyle="--", alpha=0.5)
-    ax.legend(loc="upper right")
-    ax.figure.tight_layout()
+        positionsEBD = EBD_curve.index.to_numpy()
+        velocitiesEBD = EBD_curve["Velocity [m/s]"].to_numpy()
 
-    plt.show()
+        t_be = T_be * Kt_int
+        T_berem = max(t_be - T_traction, 0)
 
+        positionsEBI = []
+        velocitiesEBI = []
 
-def computeCeilingSpeedLimits(V_permitted_mps):
+        for pos, vel in zip(positionsEBD, velocitiesEBD):
 
-    dV_ebi_min = 7.5/3.6
-    dV_ebi_max = 15.0/3.6
-    V_ebi_min = 110.0/3.6
-    V_ebi_max = 210.0/3.6
+            A_est1 = 0.1  # todo
+            A_est2 = min(A_est1, 0.4)
 
-    C_ebi = (dV_ebi_max - dV_ebi_min) / (V_ebi_max - V_ebi_min)
+            V_est = (vel - A_est1 * T_traction - A_est2 * T_berem) / (1 + v_uncertainty)
+            V_est = max(V_est, 0.0)
 
-    if V_permitted_mps <= V_ebi_min:
+            if V_est < targetVelocity:
+                velocitiesEBI.append(targetVelocity)
+                positionsEBI.append(positionsEBI[-1] + 1)
 
-        dV_ebi = dV_ebi_min
+                break
 
-    else:
-        dV_ebi = min(dV_ebi_min + C_ebi * (V_permitted_mps - V_ebi_min),dV_ebi_max,)
+            velocitiesEBI.append(V_est)
 
-    dV_warning = 0.5 * dV_ebi
-    dV_sbi = 0.75 * dV_ebi
+            V_delta_0 = V_est * v_uncertainty
+            V_delta1 = A_est1 * T_traction
+            V_delta2 = A_est2 * T_berem
+            D_bec = T_traction * (V_est + V_delta_0 + 0.5 * V_delta1) + T_berem * (
+                        V_est + V_delta_0 + V_delta1 + 0.5 * V_delta2)
+            positionsEBI.append(pos - D_bec)
 
-    return {
-        "Warning [m/s]": V_permitted_mps + dV_warning,
-        "SBI [m/s]": V_permitted_mps + dV_sbi,
-        "EBI [m/s]": V_permitted_mps + dV_ebi,
-    }
+        EBI_curve = pd.DataFrame(
+            {"Velocity [m/s]": velocitiesEBI},
+            index=positionsEBI,
+        )
 
+        EBI_curve.index.name = "Position [m]"
 
-def addStartPointToCurve(curve, velocity, start_position):
+        return EBI_curve
 
-    delta_s = curve.index.to_numpy(dtype=float)[1] - curve.index.to_numpy(dtype=float)[0]
 
-    start_point = pd.DataFrame(
-        {"Velocity [m/s]": [velocity, velocity]},
-        index=[start_position, curve.index.to_numpy(dtype=float)[0] - delta_s],
-    )
-    start_point.index.name = "Position [m]"
+    def compute_SBI_curve(self, SBI1_curve, SBI2_curve):
 
-    return pd.concat([start_point, curve])
+        positionsSBI1 = SBI1_curve.index.to_numpy()
+        velocitiesSBI1 = SBI1_curve["Velocity [m/s]"].to_numpy()
 
+        positionsSBI2 = SBI2_curve.index.to_numpy()
+        velocitiesSBI2 = SBI2_curve["Velocity [m/s]"].to_numpy()
 
-def addEndPointToCurve(curve, velocity, end_position):
+        # Use only the overlapping position range
+        minPosition = max(positionsSBI1.min(), positionsSBI2.min())
+        maxPosition = min(positionsSBI1.max(), positionsSBI2.max())
 
-    delta_s = curve.index.to_numpy(dtype=float)[-1] - curve.index.to_numpy(dtype=float)[-2]
+        step = 10.0  # [m]
+        positionsSBI = np.arange(minPosition, maxPosition + step, step)
+        positionsSBI[-1] = maxPosition
 
-    end_point = pd.DataFrame(
-        {"Velocity [m/s]": [velocity, velocity]},
-        index=[curve.index.to_numpy(dtype=float)[-1] + delta_s, end_position],
-    )
-    end_point.index.name = "Position [m]"
+        velocitiesSBI1_interpol = np.interp(positionsSBI, positionsSBI1, velocitiesSBI1)
+        velocitiesSBI2_interpol = np.interp(positionsSBI, positionsSBI2, velocitiesSBI2)
 
-    return pd.concat([curve, end_point])
+        velocitiesSBI = np.minimum(velocitiesSBI1_interpol, velocitiesSBI2_interpol)
 
+        SBI_curve = pd.DataFrame(
+            {"Velocity [m/s]": velocitiesSBI},
+            index=positionsSBI,
+        )
 
-def trimCurveToMaxVelocity(curve, maxVelocity):
+        SBI_curve.index.name = "Position [m]"
 
-    velocities = curve["Velocity [m/s]"].to_numpy(dtype=float)
-    keep_mask = velocities <= maxVelocity
+        return SBI_curve
 
-    return curve[keep_mask].copy()
 
+    def postPreProcessCurves(self, curves, target):
 
-def trimCurveFromMinVelocity(curve, minVelocity):
+        permittedVelocity = target.permittedVelocity
+        start_position = target.position - self.distancePre
 
-    velocities = curve["Velocity [m/s]"].to_numpy(dtype=float)
-    keep_mask = velocities >= minVelocity
+        speedLimits = computeCeilingSpeedLimits(permittedVelocity)
 
-    return curve[keep_mask].copy()
+        curves["EBI"] = trimCurveToMaxVelocity(curves["EBI"], speedLimits["EBI [m/s]"])
+        curves["EBI"] = addStartPointToCurve(curves["EBI"], speedLimits["EBI [m/s]"], start_position)
 
+        curves["SBI"] = trimCurveToMaxVelocity(curves["SBI"], speedLimits["SBI [m/s]"])
+        curves["SBI"] = addStartPointToCurve(curves["SBI"], speedLimits["SBI [m/s]"], start_position)
 
-def trimCurveFromMinPosition(curve, minPosition):
+        curves["W"] = trimCurveToMaxVelocity(curves["W"], speedLimits["Warning [m/s]"])
+        curves["W"] = addStartPointToCurve(curves["W"], speedLimits["Warning [m/s]"], start_position)
 
-    positions = curve.index.to_numpy(dtype=float)
-    keep_mask = positions >= minPosition
+        curves["P"] = trimCurveToMaxVelocity(curves["P"], permittedVelocity)
+        curves["P"] = addStartPointToCurve(curves["P"], permittedVelocity, start_position)
 
-    return curve[keep_mask].copy()
+        curves["I"] = trimCurveToMaxVelocity(curves["I"], permittedVelocity)
 
+        curves["EBD"] = trimCurveFromMinPosition(curves["EBD"], curves["EBI"].index.to_numpy(dtype=float)[2])
 
-def postPreProcessCurves(curves, targetPosition, permittedSpeed, plottingDistance):
+        curves["SBI2"] = trimCurveFromMinPosition(curves["SBI2"], curves["EBI"].index.to_numpy(dtype=float)[2])
 
-    start_position = targetPosition - plottingDistance
+        curves["SBD"] = trimCurveFromMinPosition(curves["SBD"], curves["SBI"].index.to_numpy(dtype=float)[2])
 
-    speedLimits = computeCeilingSpeedLimits(permittedSpeed)
+        curves["SBI1"] = trimCurveFromMinPosition(curves["SBI1"], curves["SBI"].index.to_numpy(dtype=float)[2])
 
-    curves["EBI"] = trimCurveToMaxVelocity(curves["EBI"], speedLimits["EBI [m/s]"])
-    curves["EBI"] = addStartPointToCurve(curves["EBI"], speedLimits["EBI [m/s]"], start_position)
+        return curves
 
-    curves["SBI"] = trimCurveToMaxVelocity(curves["SBI"], speedLimits["SBI [m/s]"])
-    curves["SBI"] = addStartPointToCurve(curves["SBI"], speedLimits["SBI [m/s]"], start_position)
+    def postPostProcessCurves(self, curves, target):
 
-    curves["W"] = trimCurveToMaxVelocity(curves["W"], speedLimits["Warning [m/s]"])
-    curves["W"] = addStartPointToCurve(curves["W"], speedLimits["Warning [m/s]"], start_position)
+        targetVelocity = target.targetVelocity
+        end_position = target.position + self.distancePost
 
-    curves["P"] = trimCurveToMaxVelocity(curves["P"], permittedSpeed)
-    curves["P"] = addStartPointToCurve(curves["P"], permittedSpeed, start_position)
+        speedLimits = computeCeilingSpeedLimits(targetVelocity)
 
-    curves["I"] = trimCurveToMaxVelocity(curves["I"], permittedSpeed)
+        curves["EBI"] = trimCurveFromMinVelocity(curves["EBI"], speedLimits["EBI [m/s]"])
+        curves["EBI"] = addEndPointToCurve(curves["EBI"], speedLimits["EBI [m/s]"], end_position)
 
+        curves["SBI"] = trimCurveFromMinVelocity(curves["SBI"], speedLimits["SBI [m/s]"])
+        curves["SBI"] = addEndPointToCurve(curves["SBI"], speedLimits["SBI [m/s]"], end_position)
 
-    curves["EBD"] = trimCurveFromMinPosition(curves["EBD"], curves["EBI"].index.to_numpy(dtype=float)[2])
+        curves["W"] = trimCurveFromMinVelocity(curves["W"], speedLimits["Warning [m/s]"])
+        curves["W"] = addEndPointToCurve(curves["W"], speedLimits["Warning [m/s]"], end_position)
 
-    curves["SBI2"] = trimCurveFromMinPosition(curves["SBI2"], curves["EBI"].index.to_numpy(dtype=float)[2])
+        curves["P"] = addEndPointToCurve(curves["P"], targetVelocity, end_position)
 
-    curves["SBD"] = trimCurveFromMinPosition(curves["SBD"], curves["SBI"].index.to_numpy(dtype=float)[2])
+        curves["I"] = addEndPointToCurve(curves["I"], targetVelocity, curves["P"].index.to_numpy(dtype=float)[-4])
 
-    curves["SBI1"] = trimCurveFromMinPosition(curves["SBI1"], curves["SBI"].index.to_numpy(dtype=float)[2])
+        curves["EBD"] = trimCurveFromMinVelocity(curves["EBD"], speedLimits["EBI [m/s]"])
 
-    return curves
+        curves["SBI2"] = trimCurveFromMinVelocity(curves["SBI2"], speedLimits["SBI [m/s]"])
 
+        curves["SBD"] = trimCurveFromMinVelocity(curves["SBD"], speedLimits["EBI [m/s]"])
 
-def postPostProcessCurves(curves, targetPosition, targetSpeed, postPlottingDistance):
+        curves["SBI1"] = trimCurveFromMinVelocity(curves["SBI1"], speedLimits["SBI [m/s]"])
 
-    end_position = targetPosition + postPlottingDistance
+        return curves
 
-    speedLimits = computeCeilingSpeedLimits(targetSpeed)
 
-    curves["EBI"] = trimCurveFromMinVelocity(curves["EBI"], speedLimits["EBI [m/s]"])
-    curves["EBI"] = addEndPointToCurve(curves["EBI"], speedLimits["EBI [m/s]"], end_position)
+    def computeTarget(self, target):
 
-    curves["SBI"] = trimCurveFromMinVelocity(curves["SBI"], speedLimits["SBI [m/s]"])
-    curves["SBI"] = addEndPointToCurve(curves["SBI"], speedLimits["SBI [m/s]"], end_position)
+        self.validateInput(target)
+        trainBrakingData = self.trainBrakingData
 
-    curves["W"] = trimCurveFromMinVelocity(curves["W"], speedLimits["Warning [m/s]"])
-    curves["W"] = addEndPointToCurve(curves["W"], speedLimits["Warning [m/s]"], end_position)
+        df_A_brake_safe = self.compute_A_brake_safe()
+        T_indication = max(0.8 * trainBrakingData["T_bs [s]"], 5) + self.T_driver
 
-    curves["P"] = addEndPointToCurve(curves["P"], targetSpeed, end_position)
+        curves = {}
 
-    curves["I"] = addEndPointToCurve(curves["I"], targetSpeed, curves["P"].index.to_numpy(dtype=float)[-4])
+        curves["EBD"] = self.computeBrakingCurve(df_A_brake_safe, target.SvL, target.permittedVelocity, target.targetVelocity)
 
+        curves["EBI"] = self.compute_EBI_curve(curves["EBD"], target.targetVelocity)
 
-    curves["EBD"] = trimCurveFromMinVelocity(curves["EBD"], speedLimits["EBI [m/s]"])
+        curves["SBI2"] = shiftCurveByTime(curves["EBI"], trainBrakingData["T_bs2 [s]"])
 
-    curves["SBI2"] = trimCurveFromMinVelocity(curves["SBI2"], speedLimits["SBI [m/s]"])
+        curves["SBD"] = self.computeBrakingCurve(self.trainBrakingData["A_brake_service [m/s^2]"], target.EoA, target.permittedVelocity, target.targetVelocity)
 
-    curves["SBD"] = trimCurveFromMinVelocity(curves["SBD"], speedLimits["EBI [m/s]"])
+        curves["SBI1"] = shiftCurveByTime(curves["SBD"], trainBrakingData["T_bs1 [s]"])
 
-    curves["SBI1"] = trimCurveFromMinVelocity(curves["SBI1"], speedLimits["SBI [m/s]"])
+        curves["SBI"] = self.compute_SBI_curve(curves["SBI1"], curves["SBI2"])
 
-    return curves
+        curves["W"] = shiftCurveByTime(curves["SBI"], self.T_warning)
 
+        curves["P"] = shiftCurveByTime(curves["SBI"], self.T_driver)
 
-def computeStepApproximation(curve, step=50):
+        curves["I"] = shiftCurveByTime(curves["P"], T_indication)
 
-    positions = curve.index.to_numpy(dtype=float)
-    velocities = curve["Velocity [m/s]"].to_numpy(dtype=float)
+        curves = self.postPreProcessCurves(curves, target)
 
-    samplePositions = np.arange(positions[1], positions[-1], step)
-    sampleVelocities = np.interp(samplePositions, positions, velocities)
+        if target.targetVelocity > 0:
 
-    approximation = pd.DataFrame(
-        {"Velocity [m/s]": sampleVelocities},
-        index=samplePositions,
-    )
+            curves = self.postPostProcessCurves(curves, target)
 
-    approximation.index.name = "Position [m]"
+        return curves
 
-    return approximation
 
+    def plotCurves(self, curves, target):
 
-def plotApproximation(approximation, curve, name):
+        targetPosition = target.EoA
+        permittedVelocity = target.permittedVelocity
+        targetVelocity = target.targetVelocity
 
-    curveStyles = getCurveStyles()
+        fig, ax = plt.subplots(figsize=(16, 8))
 
-    fig, ax = plt.subplots(figsize=(16, 8))
+        ax.step(
+            np.array(
+                [targetPosition - self.distancePre, targetPosition, targetPosition + self.distancePost]) / 1000,
+            np.array([permittedVelocity, permittedVelocity, targetVelocity]) * 3.6,
+            label="Speed limit", color="black", linewidth=2.0
+        )
 
-    style = {}
+        for name, curve in curves.items():
 
-    if curveStyles is not None and name in curveStyles:
-        style = curveStyles[name]
+            style = {}
 
-    ax.plot(curve.index.values / 1000, curve["Velocity [m/s]"] * 3.6, label=name, **style)
-    ax.step(approximation.index.values / 1000, approximation["Velocity [m/s]"] * 3.6, label="Step Approxmiation")
+            if self.curveStyles is not None and name in self.curveStyles:
+                style = self.curveStyles[name]
 
-    ax.set_title("ETCS Braking Curves")
-    ax.set_xlabel("Position [km]")
-    ax.set_ylabel("Velocity [km/h]")
-    ax.grid(True, which="both", linestyle="--", alpha=0.5)
-    ax.legend(loc="upper right")
-    ax.figure.tight_layout()
+            ax.plot(curve.index.values / 1000, curve["Velocity [m/s]"] * 3.6, label=name, **style)
 
-    plt.show()
+        ax.set_title("ETCS Braking Curves")
+        ax.set_xlabel("Position [km]")
+        ax.set_ylabel("Velocity [km/h]")
+        ax.grid(True, which="both", linestyle="--", alpha=0.5)
+        ax.legend(loc="upper right")
+        ax.figure.tight_layout()
 
+        plt.show()
 
-if __name__ == '__main__':
 
-    # ETCS Constants
-    T_warning = 2  # [s]
-    T_driver = 4  # [s]
+# def computeStepApproximation(curve, step=50):
+#
+#     positions = curve.index.to_numpy(dtype=float)
+#     velocities = curve["Velocity [m/s]"].to_numpy(dtype=float)
+#
+#     samplePositions = np.arange(positions[1], positions[-1], step)
+#     sampleVelocities = np.interp(samplePositions, positions, velocities)
+#
+#     approximation = pd.DataFrame(
+#         {"Velocity [m/s]": sampleVelocities},
+#         index=samplePositions,
+#     )
+#
+#     approximation.index.name = "Position [m]"
+#
+#     return approximation
+#
+#
+# def plotApproximation(approximation, curve, name):
+#
+#     curveStyles = getCurveStyles()
+#
+#     fig, ax = plt.subplots(figsize=(16, 8))
+#
+#     style = {}
+#
+#     if curveStyles is not None and name in curveStyles:
+#         style = curveStyles[name]
+#
+#     ax.plot(curve.index.values / 1000, curve["Velocity [m/s]"] * 3.6, label=name, **style)
+#     ax.step(approximation.index.values / 1000, approximation["Velocity [m/s]"] * 3.6, label="Step Approxmiation")
+#
+#     ax.set_title("ETCS Braking Curves")
+#     ax.set_xlabel("Position [km]")
+#     ax.set_ylabel("Velocity [km/h]")
+#     ax.grid(True, which="both", linestyle="--", alpha=0.5)
+#     ax.legend(loc="upper right")
+#     ax.figure.tight_layout()
+#
+#     plt.show()
+
+
+if __name__ == '__main__':
 
     # Convention:
     # Braking accelerations are stored as negative values.
@@ -461,55 +536,17 @@ def plotApproximation(approximation, curve, name):
 
     track = Track(config={'id': 'CH_StGallen_Wil'}, pathJSON='../tracks')
 
-    targetPosition = 5000  # [m]
-    overlap = 100  # [m]
-    permittedSpeed = 160 / 3.6  # [m/s]
-    targetSpeed = 24  # [m/s]
-
-    prepPlottingDistance = 3000  # [m]
-    postPlottingDistance = 1000  # [m]
-
-    SvL = targetPosition + overlap
-    EoA = targetPosition
-
-    assert permittedSpeed >= 0 and permittedSpeed < 400 / 3.6
-
-    assert targetPosition > 0 and targetPosition < track.length
-
-    assert targetSpeed >= 0 and targetSpeed < permittedSpeed
-
-
-    df_A_brake_safe = compute_A_brake_safe(trainBrakingData)
-
-    T_indication = max(0.8 * trainBrakingData["T_bs [s]"], 5) + T_driver
-
-    curves = {}
-
-    curves["EBD"] = compute_braking_curve(df_A_brake_safe, track.gradients, SvL, permittedSpeed, targetSpeed)
-
-    curves["EBI"] = compute_EBI_curve(curves["EBD"], trainBrakingData, targetSpeed)
-
-    curves["SBI2"] = shift_curve_by_time(curves["EBI"], trainBrakingData["T_bs2 [s]"])
-
-    curves["SBD"] = compute_braking_curve(trainBrakingData["A_brake_service [m/s^2]"], track.gradients, EoA, permittedSpeed, targetSpeed)
-
-    curves["SBI1"] = shift_curve_by_time(curves["SBD"], trainBrakingData["T_bs1 [s]"])
-
-    curves["SBI"] = compute_SBI_curve(curves["SBI1"], curves["SBI2"])
-
-    curves["W"] = shift_curve_by_time(curves["SBI"], T_warning)
-
-    curves["P"] = shift_curve_by_time(curves["SBI"], T_driver)
-
-    curves["I"] = shift_curve_by_time(curves["P"], T_indication)
-
-    curves = postPreProcessCurves(curves, targetPosition, permittedSpeed, prepPlottingDistance)
-
-    if targetSpeed > 0:
+    target = BrakingTarget(
+            position=5000,
+            overlap= 100,
+            permittedVelocity=160/3.6,
+            targetVelocity=0/3.6
+    )
 
-        curves = postPostProcessCurves(curves, targetPosition, targetSpeed, postPlottingDistance)
+    calculator = EtcsBrakingCurveCalculator(trainBrakingData, track, distancePre=5000, distancePost=1000)
+    curve_set = calculator.computeTarget(target)
 
-    plotCurves(curves, targetPosition, permittedSpeed, targetSpeed, prepPlottingDistance, postPlottingDistance)
+    calculator.plotCurves(curve_set, target)
 
     # approximation = computeStepApproximation(curves["P"])
     # plotApproximation(approximation, curves["P"], "P")

From ccf8da74bf07b6e11f94288b3096b0cffb9fcd02 Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Mon, 8 Jun 2026 13:28:21 +0200
Subject: [PATCH 30/31] revisions pull request

---
 mseetc/etcs.py  |  47 +----------
 mseetc/track.py | 216 ++++++++++++++++++++++--------------------------
 mseetc/train.py |   6 +-
 3 files changed, 103 insertions(+), 166 deletions(-)

diff --git a/mseetc/etcs.py b/mseetc/etcs.py
index 1b1e457..7368336 100644
--- a/mseetc/etcs.py
+++ b/mseetc/etcs.py
@@ -465,48 +465,6 @@ def plotCurves(self, curves, target):
         plt.show()
 
 
-# def computeStepApproximation(curve, step=50):
-#
-#     positions = curve.index.to_numpy(dtype=float)
-#     velocities = curve["Velocity [m/s]"].to_numpy(dtype=float)
-#
-#     samplePositions = np.arange(positions[1], positions[-1], step)
-#     sampleVelocities = np.interp(samplePositions, positions, velocities)
-#
-#     approximation = pd.DataFrame(
-#         {"Velocity [m/s]": sampleVelocities},
-#         index=samplePositions,
-#     )
-#
-#     approximation.index.name = "Position [m]"
-#
-#     return approximation
-#
-#
-# def plotApproximation(approximation, curve, name):
-#
-#     curveStyles = getCurveStyles()
-#
-#     fig, ax = plt.subplots(figsize=(16, 8))
-#
-#     style = {}
-#
-#     if curveStyles is not None and name in curveStyles:
-#         style = curveStyles[name]
-#
-#     ax.plot(curve.index.values / 1000, curve["Velocity [m/s]"] * 3.6, label=name, **style)
-#     ax.step(approximation.index.values / 1000, approximation["Velocity [m/s]"] * 3.6, label="Step Approxmiation")
-#
-#     ax.set_title("ETCS Braking Curves")
-#     ax.set_xlabel("Position [km]")
-#     ax.set_ylabel("Velocity [km/h]")
-#     ax.grid(True, which="both", linestyle="--", alpha=0.5)
-#     ax.legend(loc="upper right")
-#     ax.figure.tight_layout()
-#
-#     plt.show()
-
-
 if __name__ == '__main__':
 
     # Convention:
@@ -546,7 +504,4 @@ def plotCurves(self, curves, target):
     calculator = EtcsBrakingCurveCalculator(trainBrakingData, track, distancePre=5000, distancePost=1000)
     curve_set = calculator.computeTarget(target)
 
-    calculator.plotCurves(curve_set, target)
-
-    # approximation = computeStepApproximation(curves["P"])
-    # plotApproximation(approximation, curves["P"], "P")
+    calculator.plotCurves(curve_set, target)
\ No newline at end of file
diff --git a/mseetc/track.py b/mseetc/track.py
index b683402..086c2d0 100644
--- a/mseetc/track.py
+++ b/mseetc/track.py
@@ -116,72 +116,85 @@ def computeDiscretizationPoints(track, numIntervals, opts):
     return df3
 
 
-def adaptConstantTrackAttributesToNewShootingNodes(df3, numIntervals):
+def adaptConstantTrackAttributesToNewShootingNodes(df, numIntervals):
+    '''
+    Adapt gradient and curvature values to the current shooting nodes.
 
-    if "Gradient linear term [permil/m]" in df3.columns:
+    For piecewise linear profiles, the constant term at each node is updated using the linear term from the previous interval.
 
-        positions = df3.index.to_numpy(dtype=float)
-        grads = [df3["Gradient [permil]"].iloc[0]]
+    Example:
+        Gradient at 0 m is 10 permil and the linear term is 0.01 permil/m.
+        At a new shooting node at 100 m, the gradient becomes:
+
+            10 + 100 * 0.01 = 11 permil
+
+        Without a linear term, the value would remain 10 permil.
+    '''
+
+    if "Gradient linear term [permil/m]" in df.columns:
+
+        positions = df.index.to_numpy(dtype=float)
+        grads = [df["Gradient [permil]"].iloc[0]]
 
         for idx in range(1, numIntervals + 1):
 
-            if np.isclose(df3["Gradient [permil]"].iloc[idx - 1], df3["Gradient [permil]"].iloc[idx]):
+            if np.isclose(df["Gradient [permil]"].iloc[idx - 1], df["Gradient [permil]"].iloc[idx]):
 
-                grads.append(grads[-1] + (positions[idx] - positions[idx - 1]) * df3["Gradient linear term [permil/m]"].iloc[idx - 1])
+                grads.append(grads[-1] + (positions[idx] - positions[idx - 1]) * df["Gradient linear term [permil/m]"].iloc[idx - 1])
 
             else:
 
-                grads.append(df3["Gradient [permil]"].iloc[idx])
+                grads.append(df["Gradient [permil]"].iloc[idx])
 
-        df3["Gradient [permil]"] = grads
+        df["Gradient [permil]"] = grads
 
     else:
 
-        df3["Gradient linear term [permil/m]"] = np.zeros(len(df3))
+        df["Gradient linear term [permil/m]"] = np.zeros(len(df))
 
 
-    if "Curvature linear term [1/m^2]" in df3.columns:
+    if "Curvature linear term [1/m^2]" in df.columns:
 
-        positions = df3.index.to_numpy(dtype=float)
-        curvs = [df3["Curvature [1/m]"].iloc[0]]
+        positions = df.index.to_numpy(dtype=float)
+        curvs = [df["Curvature [1/m]"].iloc[0]]
 
         for idx in range(1, numIntervals + 1):
 
-            if np.isclose(df3["Curvature [1/m]"].iloc[idx - 1], df3["Curvature [1/m]"].iloc[idx]):
+            if np.isclose(df["Curvature [1/m]"].iloc[idx - 1], df["Curvature [1/m]"].iloc[idx]):
 
-                curvs.append(curvs[-1] + (positions[idx] - positions[idx - 1]) * df3["Curvature linear term [1/m^2]"].iloc[idx - 1])
+                curvs.append(curvs[-1] + (positions[idx] - positions[idx - 1]) * df["Curvature linear term [1/m^2]"].iloc[idx - 1])
 
             else:
 
-                curvs.append(df3["Curvature [1/m]"].iloc[idx])
+                curvs.append(df["Curvature [1/m]"].iloc[idx])
 
-        df3["Curvature [1/m]"] = curvs
+        df["Curvature [1/m]"] = curvs
 
     else:
 
-        df3["Curvature linear term [1/m^2]"] = np.zeros(len(df3))
+        df["Curvature linear term [1/m^2]"] = np.zeros(len(df))
 
 
-def makePwcLengthDependentTrackAttibutes(df3):
+def makePwcLengthDependentTrackAttibutes(df):
 
-    positions = df3.index.to_numpy(dtype=float)
+    positions = df.index.to_numpy(dtype=float)
 
-    g_pwl = df3["Gradient [permil]"].to_numpy(dtype=float)
-    g_linear = df3["Gradient linear term [permil/m]"].to_numpy(dtype=float)
+    g_pwl = df["Gradient [permil]"].to_numpy(dtype=float)
+    g_linear = df["Gradient linear term [permil/m]"].to_numpy(dtype=float)
 
-    c_pwl = df3["Curvature [1/m]"].to_numpy(dtype=float)
-    c_linear = df3["Curvature linear term [1/m^2]"].to_numpy(dtype=float)
+    c_pwl = df["Curvature [1/m]"].to_numpy(dtype=float)
+    c_linear = df["Curvature linear term [1/m^2]"].to_numpy(dtype=float)
 
     ds = positions[1:] - positions[:-1]
 
     g_pwc = g_pwl[:-1] + 0.5 * g_linear[:-1] * ds
     c_pwc = c_pwl[:-1] + 0.5 * c_linear[:-1] * ds
 
-    df3["Gradient [permil]"] = np.r_[g_pwc, g_pwc[-1]]
-    df3["Curvature [1/m]"] = np.r_[c_pwc, c_pwc[-1]]
+    df["Gradient [permil]"] = np.r_[g_pwc, g_pwc[-1]]
+    df["Curvature [1/m]"] = np.r_[c_pwc, c_pwc[-1]]
 
-    df3["Gradient linear term [permil/m]"] = np.zeros(len(df3))
-    df3["Curvature linear term [1/m^2]"] = np.zeros(len(df3))
+    df["Gradient linear term [permil/m]"] = np.zeros(len(df))
+    df["Curvature linear term [1/m^2]"] = np.zeros(len(df))
 
 
 class Track():
@@ -631,100 +644,51 @@ def updateSpeedLimitsToTrainLength(self, trainLength):
             )
 
 
-    def updateGradientsToTrainLength(self, trainLength):
+    def updateTrackAttributeToTrainLength(self, df, trainLength):
         """
-        Convert pointwise gradient changes into train-length-dependent piecewise linear gradients.
+        Convert pointwise track attribute changes into train-length-dependent piecewise linear values.
 
-        A gradient step at position s does not affect the whole train at once.
-        Instead, its effect is spread over one train length, from s to s + trainLength.
+        A step at position s does not affect the whole train at once. Instead, its
+        effect is spread over one train length, from s to s + trainLength.
         """
 
-        gradientValues = self.gradients["Gradient [permil]"].to_numpy(dtype=float)
-        gradientPositions = self.gradients.index.to_numpy(dtype=float)
-
-        if len(gradientPositions) <= 1:
-
-            return
-
-        # assume train starts on a flat track
-        gradientValues = np.r_[0, gradientValues]
-        gradientPositions = np.r_[-trainLength, gradientPositions]
-
-        # Each original gradient jump is spread linearly over one train length assuming uniform mass.
-        # The first point has no previous gradient, so its slope contribution is zero.
-        gradientJumpSlopes = np.r_[0.0, (gradientValues[1:] - gradientValues[:-1]) / trainLength]
-
-        # New breakpoints occur both when the front of the train reaches a gradient and when the rear of the train has passed it.
-        adjustedPositions = np.sort(np.unique(np.r_[gradientPositions, gradientPositions + trainLength]))
-        adjustedPositions = adjustedPositions[adjustedPositions < self.length]
-
-        adjustedGradients = [0]
-        gradientLinearTerms = [gradientValues[0]/trainLength]
-
-        epsilon = 1e-3
-
-        for idx in range(1,len(adjustedPositions)):
-
-            currentPosition = adjustedPositions[idx]
-            previousPosition = adjustedPositions[idx - 1]
-
-            intervalLength = currentPosition - previousPosition
-
-            # Continue the previous linear gradient to the current position.
-            currentGradient = adjustedGradients[-1] + intervalLength * gradientLinearTerms[-1]
-
-            # Active gradient jumps are those currently within one train length behind the train front.
-            activeJumpMask = (
-                    (currentPosition - trainLength + epsilon < gradientPositions)
-                    & (gradientPositions < currentPosition + epsilon)
-            )
-
-            currentLinearTerm = np.sum(gradientJumpSlopes[activeJumpMask])
+        if "Gradient [permil]" in df.columns:
 
-            adjustedGradients.append(currentGradient)
-            gradientLinearTerms.append(currentLinearTerm)
+            valueColumn = "Gradient [permil]"
+            linearTermColumn = "Gradient linear term [permil/m]"
+            plotFunction = plotGradients
 
-        adjustedPositions = adjustedPositions[1:]
-        adjustedGradients = adjustedGradients[1:]
-        gradientLinearTerms = gradientLinearTerms[1:]
-
-        if plotDebug:
-
-            plotGradients(self, np.asarray(adjustedPositions, dtype=float), np.asarray(adjustedGradients, dtype=float), np.asarray(gradientLinearTerms, dtype=float))
+        elif "Curvature [1/m]" in df.columns:
 
-        self.gradients = pd.DataFrame({"Gradient [permil]": adjustedGradients, "Gradient linear term [permil/m]": gradientLinearTerms}, index=adjustedPositions)
+            valueColumn = "Curvature [1/m]"
+            linearTermColumn = "Curvature linear term [1/m^2]"
+            plotFunction = plotCurvatures
 
+        else:
 
-    def updateCurvaturesToTrainLength(self, trainLength):
-        """
-        Convert pointwise curvature changes into train-length-dependent piecewise linear curvatures.
-
-        A curvature step at position s does not affect the whole train at once.
-        Instead, its effect is spread over one train length, from s to s + trainLength.
-        """
+            raise ValueError("Unknown track attribute DataFrame!")
 
-        curvatureValues = self.curvatures["Curvature [1/m]"].to_numpy(dtype=float)
-        curvaturePositions = self.curvatures.index.to_numpy(dtype=float)
+        values = df[valueColumn].to_numpy(dtype=float)
+        positions = df.index.to_numpy(dtype=float)
 
-        if len(curvaturePositions) <= 1:
+        if len(positions) <= 1:
 
-            return
+            return df
 
-        # assume train starts on a straight track
-        curvatureValues = np.r_[0, curvatureValues]
-        curvaturePositions = np.r_[-trainLength, curvaturePositions]
+        # Assume the train starts before the track with a flat and straight track.
+        values = np.r_[0.0, values]
+        positions = np.r_[-trainLength, positions]
 
-        # Each original curvature jump is spread linearly over one train length assuming uniform mass.
-        # The first point has no previous curvature, so its slope contribution is zero.
-        curvatureJumpSlopes = np.r_[0.0, (curvatureValues[1:] - curvatureValues[:-1]) / trainLength]
+        # Each original step is spread linearly over one train length assuming uniform mass.
+        # The first point has no previous value, so its slope contribution is zero.
+        jumpSlopes = np.r_[0.0, (values[1:] - values[:-1]) / trainLength]
 
-        # New breakpoints occur both when the front of the train reaches a curvature
-        # change and when the rear of the train has passed it.
-        adjustedPositions = np.sort(np.unique(np.r_[curvaturePositions, curvaturePositions + trainLength]))
+        # New breakpoints occur both when the front of the train reaches a step and when the rear of the train has passed it.
+        adjustedPositions = np.sort(np.unique(np.r_[positions, positions + trainLength]))
         adjustedPositions = adjustedPositions[adjustedPositions < self.length]
 
-        adjustedCurvatures = [curvatureValues[0]]
-        curvatureLinearTerms = [0.0]
+        adjustedValues = [0.0]
+        linearTerms = [0.0]
 
         epsilon = 1e-3
 
@@ -735,29 +699,47 @@ def updateCurvaturesToTrainLength(self, trainLength):
 
             intervalLength = currentPosition - previousPosition
 
-            # Continue the previous linear curvature to the current position.
-            currentCurvature = adjustedCurvatures[-1] + intervalLength * curvatureLinearTerms[-1]
+            # Continue the previous linear value to the current position.
+            currentValue = adjustedValues[-1] + intervalLength * linearTerms[-1]
 
-            # Active curvature jumps are those currently within one train length behind the train front.
+            # Active jumps are those currently within one train length behind the train front.
             activeJumpMask = (
-                    (currentPosition - trainLength + epsilon < curvaturePositions)
-                    & (curvaturePositions < currentPosition + epsilon)
+                    (currentPosition - trainLength + epsilon < positions)
+                    & (positions < currentPosition + epsilon)
             )
 
-            currentLinearTerm = np.sum(curvatureJumpSlopes[activeJumpMask])
+            currentLinearTerm = np.sum(jumpSlopes[activeJumpMask])
 
-            adjustedCurvatures.append(currentCurvature)
-            curvatureLinearTerms.append(currentLinearTerm)
+            adjustedValues.append(currentValue)
+            linearTerms.append(currentLinearTerm)
 
+        # Remove artificial point at -trainLength.
         adjustedPositions = adjustedPositions[1:]
-        adjustedCurvatures = adjustedCurvatures[1:]
-        curvatureLinearTerms = curvatureLinearTerms[1:]
+        adjustedValues = adjustedValues[1:]
+        linearTerms = linearTerms[1:]
+
+        if plotDebug and plotFunction is not None:
+
+            plotFunction(self, np.asarray(adjustedPositions, dtype=float), np.asarray(adjustedValues, dtype=float), np.asarray(linearTerms, dtype=float)
+)
+
+        return pd.DataFrame(
+            {
+                valueColumn: adjustedValues,
+                linearTermColumn: linearTerms
+            },
+            index=adjustedPositions
+        )
 
-        if plotDebug:
 
-            plotCurvatures(self, np.asarray(adjustedPositions, dtype=float), np.asarray(adjustedCurvatures, dtype=float), np.asarray(curvatureLinearTerms, dtype=float))
+    def updateGradientsToTrainLength(self, trainLength):
+
+        self.gradients = self.updateTrackAttributeToTrainLength(self.gradients, trainLength)
+
+
+    def updateCurvaturesToTrainLength(self, trainLength):
 
-        self.curvatures = pd.DataFrame({"Curvature [1/m]": adjustedCurvatures, "Curvature linear term [1/m^2]": curvatureLinearTerms}, index=adjustedPositions)
+        self.curvatures = self.updateTrackAttributeToTrainLength(self.curvatures, trainLength)
 
 
 if __name__ == '__main__':
diff --git a/mseetc/train.py b/mseetc/train.py
index 7176ae5..9ef6d4c 100644
--- a/mseetc/train.py
+++ b/mseetc/train.py
@@ -117,7 +117,7 @@ def __init__(self, config, pathJSON=Path(__file__).parent.parent / 'trains') ->
             tunnel_data = data["tunnel resistance"] # additive aerodynamic tunnel drag as dict per tunnel cross section
 
             cross_section_unit = tunnel_data["units"]["cross section"] # tunnel cross section [m^2]
-            resistance_unit = tunnel_data["units"]["coefficient"] # resistance coefficient [kg/m]
+            resistance_unit = tunnel_data["units"]["coefficient"] # resistance coefficient [N/(m/s)^2]
 
             self.tunnelCoefficients = {
                 convertUnit(cross_section, cross_section_unit): convertUnit(
@@ -290,8 +290,8 @@ def __init__(self, sr0, sr1, sr2, rho=1, g=9.81, withPnBrake=True) -> None:
 
         rollingResistance = sr0 + sr1*ca.sqrt(velocitySquared) + sr2*velocitySquared  # [N/kg]
         gradientResistance = g*(1/rho)*(gradient+gradientLinearTerm*position)  # [N/kg]
-        curvatureResistance = (1/rho)*(5.07 * (curvature+curvatureLinearTerm*position))  # [N/kg]
-        tunnelResistance = tunnelFactor * velocitySquared  # [N/kg]
+        curvatureResistance = (1/rho)*(5.07*(curvature+curvatureLinearTerm*position))  # [N/kg]
+        tunnelResistance = tunnelFactor*velocitySquared  # [N/kg]
 
         acceleration = traction + (pnBrake if withPnBrake else 0) - rollingResistance - gradientResistance - curvatureResistance - tunnelResistance  # [m/s^2]
 

From 7c986cab001f041af084808ca038ba800e07720e Mon Sep 17 00:00:00 2001
From: Roli15 <81290196+Roli15@users.noreply.github.com>
Date: Mon, 8 Jun 2026 14:45:11 +0200
Subject: [PATCH 31/31] small refactoring of etcs.py

---
 mseetc/etcs.py | 96 +++++++++++++++++++++++++++++++-------------------
 1 file changed, 60 insertions(+), 36 deletions(-)

diff --git a/mseetc/etcs.py b/mseetc/etcs.py
index 7368336..9d56885 100644
--- a/mseetc/etcs.py
+++ b/mseetc/etcs.py
@@ -109,25 +109,39 @@ class BrakingTarget:
     permittedVelocity: float  # [m/s]
     targetVelocity: float  # [m/s]
 
+    # EoA: End of authority
     @property
     def EoA(self):
         return self.position
 
+    # SvL: Supervised location
     @property
     def SvL(self):
         return self.position + self.overlap
 
 
 class EtcsBrakingCurveCalculator:
+    """
+    Conventions
+    -----------
+    - Position increases in the train running direction.
+    - Braking accelerations are stored as negative values.
+    - Gradient is positive for uphill and negative for downhill.
+    - A_gradient = g * gradient / 1000.
+    - Curves are computed backwards from the target position.
+    """
 
     def __init__(self, trainBrakingData, track, distancePre=3000, distancePost=1000):
 
         self.trainBrakingData = trainBrakingData
         self.track = track
 
+        # Compute the braking curve using fixed time steps of length dt.
         self.dt = 0.1  # [s]
-        self.distancePre = distancePre
-        self.distancePost = distancePost
+
+        # speed decrease is plotted from BrakingTarget.position - distancePre until BrakingTarget.position + distancePost
+        self.distancePre = distancePre  # [m]
+        self.distancePost = distancePost  # [m]
 
         # ETCS Constants
         self.T_warning = 2.0  # [s]
@@ -161,7 +175,7 @@ def validateInput(self, target):
             raise ValueError("targetVelocity must be lower than permittedVelocity.")
 
 
-    def compute_A_brake_safe(self):
+    def computeABrakeSafe(self):
 
         trainBrakingData = self.trainBrakingData
 
@@ -187,7 +201,7 @@ def compute_A_brake_safe(self):
         }
 
 
-    def compute_A_gradient(self, currentPosition):
+    def computeAGradient(self, currentPosition):
 
         positions = self.track.gradients.index.values
         gradients = self.track.gradients["Gradient [permil]"].values
@@ -197,35 +211,46 @@ def compute_A_gradient(self, currentPosition):
 
         if idx == 0:
 
-            return 0  # [start of track has been exceeded]
+            # If the backward-computed curve extends before the first known gradient point, assume flat track.
+            return 0
 
         gradient = gradients[idx]
         return 9.81 * gradient * 0.001
 
 
-    def computeBrakingCurve(self, brakingProfile, target_position, permittedVelocity, targetVelocity):
+    def computeBrakingCurve(self, brakingProfile, targetPosition, permittedVelocity, targetVelocity):
+        """
+        Compute a braking curve backwards from a target position and target velocity.
+
+        The braking profile is defined by velocity thresholds and corresponding braking decelerations.
+        The curve is integrated backwards in fixed time steps until the maximum relevant velocity is reached.
+        """
 
-        max_velocity = permittedVelocity * 1.4
+        maxVelocity = permittedVelocity * 1.4
 
-        positions = [target_position]
+        positions = [targetPosition]
         velocities = [targetVelocity]
 
-        threshold_velocities = list(brakingProfile["velocity [m/s]"])
-        braking_values = list(brakingProfile["value [m/s^2]"])
+        thresholdVelocities = list(brakingProfile["velocity [m/s]"])
+        brakingValues = list(brakingProfile["value [m/s^2]"])
+
+        openEndedVelocityThreshold  = 200  # [m/s], practical upper bound for open-ended last interval of the braking profile, basically inf velocity for a train
+        thresholdVelocities.append(openEndedVelocityThreshold )
 
-        threshold_velocities.append(200)  # basically inf velocity
+        for idx in range(len(thresholdVelocities) - 1):
 
-        for idx in range(len(threshold_velocities) - 1):
+            upperThreshold = thresholdVelocities[idx + 1]
+
+            # The curve starts at targetVelocity, so lower velocity ranges are not relevant.
+            if targetVelocity > upperThreshold:
 
-            if targetVelocity > threshold_velocities[idx + 1]:
                 continue
 
-            threshold_velocity = threshold_velocities[idx + 1]
-            A_brake = braking_values[idx]
+            A_brake = brakingValues[idx]
 
-            while True:
+            while velocities[-1] < upperThreshold and velocities[-1] < maxVelocity:
 
-                A_gradient = self.compute_A_gradient(positions[-1])
+                A_gradient = self.computeAGradient(positions[-1])
 
                 v_old = velocities[-1]
                 v_new = v_old - (A_brake - A_gradient) * self.dt
@@ -234,10 +259,8 @@ def computeBrakingCurve(self, brakingProfile, target_position, permittedVelocity
                 positions.append(x_new)
                 velocities.append(v_new)
 
-                if v_new >= threshold_velocity or v_new >= max_velocity:
-                    break
+            if velocities[-1] >= maxVelocity:
 
-            if velocities[-1] >= max_velocity:
                 break
 
         curve = pd.DataFrame(
@@ -250,7 +273,7 @@ def computeBrakingCurve(self, brakingProfile, target_position, permittedVelocity
         return curve
 
 
-    def compute_EBI_curve(self, EBD_curve, targetVelocity):
+    def computeEBICurve(self, EBD_curve, targetVelocity):
 
         trainBrakingData = self.trainBrakingData
 
@@ -277,6 +300,10 @@ def compute_EBI_curve(self, EBD_curve, targetVelocity):
             V_est = max(V_est, 0.0)
 
             if V_est < targetVelocity:
+
+                # Stop once the estimated velocity drops below the target velocity.
+                # Add the target velocity as the final point, using a small position offset as a simplified position.
+                # This is only used to terminate the plotted EBI curve at targetVelocity.
                 velocitiesEBI.append(targetVelocity)
                 positionsEBI.append(positionsEBI[-1] + 1)
 
@@ -287,8 +314,7 @@ def compute_EBI_curve(self, EBD_curve, targetVelocity):
             V_delta_0 = V_est * v_uncertainty
             V_delta1 = A_est1 * T_traction
             V_delta2 = A_est2 * T_berem
-            D_bec = T_traction * (V_est + V_delta_0 + 0.5 * V_delta1) + T_berem * (
-                        V_est + V_delta_0 + V_delta1 + 0.5 * V_delta2)
+            D_bec = T_traction * (V_est + V_delta_0 + 0.5 * V_delta1) + T_berem * (V_est + V_delta_0 + V_delta1 + 0.5 * V_delta2)
             positionsEBI.append(pos - D_bec)
 
         EBI_curve = pd.DataFrame(
@@ -301,7 +327,7 @@ def compute_EBI_curve(self, EBD_curve, targetVelocity):
         return EBI_curve
 
 
-    def compute_SBI_curve(self, SBI1_curve, SBI2_curve):
+    def computeSBICurve(self, SBI1_curve, SBI2_curve):
 
         positionsSBI1 = SBI1_curve.index.to_numpy()
         velocitiesSBI1 = SBI1_curve["Velocity [m/s]"].to_numpy()
@@ -320,6 +346,8 @@ def compute_SBI_curve(self, SBI1_curve, SBI2_curve):
         velocitiesSBI1_interpol = np.interp(positionsSBI, positionsSBI1, velocitiesSBI1)
         velocitiesSBI2_interpol = np.interp(positionsSBI, positionsSBI2, velocitiesSBI2)
 
+        # At each position, take the lower speed of SBI1 and SBI2.
+        # This gives the more restrictive plotted SBI curve.
         velocitiesSBI = np.minimum(velocitiesSBI1_interpol, velocitiesSBI2_interpol)
 
         SBI_curve = pd.DataFrame(
@@ -332,7 +360,7 @@ def compute_SBI_curve(self, SBI1_curve, SBI2_curve):
         return SBI_curve
 
 
-    def postPreProcessCurves(self, curves, target):
+    def processCurvesBeforeTarget(self, curves, target):
 
         permittedVelocity = target.permittedVelocity
         start_position = target.position - self.distancePre
@@ -363,7 +391,7 @@ def postPreProcessCurves(self, curves, target):
 
         return curves
 
-    def postPostProcessCurves(self, curves, target):
+    def processCurvcesAfterTarget(self, curves, target):
 
         targetVelocity = target.targetVelocity
         end_position = target.position + self.distancePost
@@ -399,14 +427,14 @@ def computeTarget(self, target):
         self.validateInput(target)
         trainBrakingData = self.trainBrakingData
 
-        df_A_brake_safe = self.compute_A_brake_safe()
+        ABrakeSafeProfile = self.computeABrakeSafe()
         T_indication = max(0.8 * trainBrakingData["T_bs [s]"], 5) + self.T_driver
 
         curves = {}
 
-        curves["EBD"] = self.computeBrakingCurve(df_A_brake_safe, target.SvL, target.permittedVelocity, target.targetVelocity)
+        curves["EBD"] = self.computeBrakingCurve(ABrakeSafeProfile, target.SvL, target.permittedVelocity, target.targetVelocity)
 
-        curves["EBI"] = self.compute_EBI_curve(curves["EBD"], target.targetVelocity)
+        curves["EBI"] = self.computeEBICurve(curves["EBD"], target.targetVelocity)
 
         curves["SBI2"] = shiftCurveByTime(curves["EBI"], trainBrakingData["T_bs2 [s]"])
 
@@ -414,7 +442,7 @@ def computeTarget(self, target):
 
         curves["SBI1"] = shiftCurveByTime(curves["SBD"], trainBrakingData["T_bs1 [s]"])
 
-        curves["SBI"] = self.compute_SBI_curve(curves["SBI1"], curves["SBI2"])
+        curves["SBI"] = self.computeSBICurve(curves["SBI1"], curves["SBI2"])
 
         curves["W"] = shiftCurveByTime(curves["SBI"], self.T_warning)
 
@@ -422,11 +450,11 @@ def computeTarget(self, target):
 
         curves["I"] = shiftCurveByTime(curves["P"], T_indication)
 
-        curves = self.postPreProcessCurves(curves, target)
+        curves = self.processCurvesBeforeTarget(curves, target)
 
         if target.targetVelocity > 0:
 
-            curves = self.postPostProcessCurves(curves, target)
+            curves = self.processCurvcesAfterTarget(curves, target)
 
         return curves
 
@@ -467,10 +495,6 @@ def plotCurves(self, curves, target):
 
 if __name__ == '__main__':
 
-    # Convention:
-    # Braking accelerations are stored as negative values.
-    # Gradient acceleration is positive for uphill and negative for downhill.
-
     trainBrakingData = {
         "A_brake_emergency [m/s^2]": {
             "velocity [m/s]": [0, 20, 40, 60],