# Chained Quantile Morphing with Normalizing Flows

Samuel Bright-Thonney<EMAIL_ADDRESS>Physics Department, Cornell
University, 109 Clark Hall, Ithaca, New York 14853, USA  Philip Harris
<EMAIL_ADDRESS>Physics Department, Massachusetts Institute of Technology,
Cambridge, MA 02139, USA  The NSF AI Institute for Artificial Intelligence and
Fundamental Interactions, USA  Patrick McCormack<EMAIL_ADDRESS>Physics
Department, Massachusetts Institute of Technology, Cambridge, MA 02139, USA
The NSF AI Institute for Artificial Intelligence and Fundamental Interactions,
USA  Simon Rothman<EMAIL_ADDRESS>Physics Department, Massachusetts
Institute of Technology, Cambridge, MA 02139, USA  The NSF AI Institute for
Artificial Intelligence and Fundamental Interactions, USA

###### Abstract

Accounting for inaccuracies in Monte Carlo simulations is a crucial step in
any high energy physics analysis. It becomes especially important when
training machine learning models, which can amplify simulation inaccuracies
and introduce large discrepancies and systematic uncertainties when the model
is applied to data. In this paper, we introduce a method to transform
simulated events to better match data using normalizing flows, a class of deep
learning-based density estimation models. Our proposal uses a technique called
chained quantile morphing, which corrects a set of observables by iteratively
shifting each entry according to a conditonal cumulative density function. We
demonstrate the technique on a realistic particle physics dataset, and compare
it to a neural network-based reweighting method. We also introduce a new
contrastive learning technique to correct high dimensional particle-level
inputs, which naively cannot be efficiently corrected with morphing
strategies.

##  I Introduction

Searches and measurements using Large Hadron Collider (LHC) data almost always
rely on Monte Carlo (MC) simulations to develop analysis, validate tools, and
frequently predict backgrounds. These simulations are widely acknowledged to
be imperfect, particularly in modeling detector interactions and the non-
perturbative physics of hadronization; data-driven strategies are preferred
when possible. Despite the limitations of MC, many modern analyses rely
heavily on machine learning (ML) to maximize their sensitivity, and these
algorithms are typically trained on MC. The use of ML allows for the effective
utilization of complex patterns and correlations in high-dimensional data and
is thus more powerful, but also particularly susceptible to spurious,
unphysical artifacts present in the simulations.

This issue has only grown in recent years, with a community-wide move towards
training ML models on extremely granular, particle-level information with
architectures such as ParticleNet Qu and Gouskos (2020); CMS (2023), LundNet
Dreyer and Qu (2021); ATL (2023), and the Dynamic Reduction Network Rothman
(2021). As the reliance on finer details increases, simulations become less
reliable. These inaccuracies can lead to significant discrepancies between a
model’s performance on MC and real experimental data. This adds additional
work for physicists (deriving corrections, scale factors, etc.), introduces
new uncertainties, and points to the deeper issue of training our most
powerful analysis tools on flawed simulations. It is conceivable that, in the
near future, ML-related systematic uncertainties coming from flaws in the
simulation will be the primary limitation on the precision of Standard Model
measurements or the sensitivity of new physics searches.

In this paper, we introduce a general purpose strategy to transform samples
from one probability distribution to match another using a deep learning
implementation of chained quantile morphing (CQM) Spyromitros-Xioufis  _et
al._ (2016); Collaboration (2022). CQM iteratively transforms a set of
observables $\mathbf{x}=(x_{1},\ldots,x_{N})$ using the conditional cumulative
distribution functions (CDFs) $F_{i}^{\mathrm{MC}}(x_{i}|\mathbf{x}_{1:i-1})$
and $F_{i}^{\mathrm{Data}}(x_{i}|\mathbf{x}_{1:i-1})$, and was first used for
LHC analysis to improve the quality of photon identification in the Compact
Muon Solenoid detector Ref. Collaboration (2022). The authors used discretized
approximations of the CDFs to correct MC inaccuracies, which allowed them to
reduce an important systematic uncertainty. In this work, we develop a
continuous and precise version of their approach using normalizing flows
Kobyzev  _et al._ (2021); Tabak and Turner ; Papamakarios  _et al._ (2021) – a
family of invertible ML models capable of learning complex conditional
probability densities. While CQM can be used to transform between any two
distributions of the same dimensionality, we focus on the high energy physics
context of transforming simulated Monte Carlo observables to better match
experimental data.

Monte Carlo correction strategies typically fall into two categories:
reweighting Cranmer  _et al._ (2015); Andreassen and Nachman (2020);
Andreassen  _et al._ (2021); Nachman and Thaler (2020, 2022); Andreassen  _et
al._ (2020) and morphing Klein  _et al._ (2022); Golling  _et al._ (2023a, b);
Raine  _et al._ (2023); Sengupta  _et al._ (2023); Algren  _et al._ (2023).
CQM is a morphing strategy, meaning we correct the values of observables
$\mathbf{x}_{\mathrm{MC}}\to\mathbf{x}_{\mathrm{MC}}^{\mathrm{corr}}$ in a way
that results in the overall distribution agreeing better with data.
Reweighting methods learn a per-event weight that improves data/MC agreement
without explicitly altering any observables. Both are effective strategies and
have been extensively studied. However, the underlying motivation for each of
these methodologies is different due to the nature of how they correct MC.
Morphing methods will shift full distributions, thereby breaking relations of
parameters within the simulation. Such a correction is applicable when
considering a recalibration of a detector readout where the full distribution
is shifted to correct an mis-modeled relation between a generated effect and a
reconstructed effect. An example from LHC physics could be photon energy
corrections, whereby one shifts the mis-reconstructed energy spectrum to match
the generated spectrum. Reweighting strategies are often used when there is a
need to preserve invariant quantities such as particle mass. An example of
such a strategy is reweighting simulated top quark momentum spectra to match
the observed data distribution.

While both reweighting and morphing have been shown to be effective, when
distributions differ by large amounts, reweighting strategies can lead to
large uncertainties in their prediction due to the presence of high weights.
As a result, morphing may be more effective, particularly when there are
significant differences between data and MC111One clear example is a
discrepancy in the tails of data and MC distributions, where an event weight
cannot make up for lack of events. Morphing also produces a new dataset which
may be easier to use in downstream applications such as training ML models.

In the following paper, we develop and demonstrate the effective use of flow-
based CQM. We discuss the details of the flow-based implementation of CQM in
Sec. II and contrast it with existing approaches. In Sec. III we present
results using CQM to morph between a pair of toy 2D distributions and a pair
of realistic simulated particle physics datasets. Sec. IV explores the
possibility of applying CQM in very high-dimensional (e.g. particle-level)
contexts by embedding the physically relevant information into a low-
dimensional space using contrastive learning. Finally, in Sec. V we
demonstrate that CQM is insensitive to small levels of signal contamination
and can be trained in a control region and interpolated accurately into a
blinded signal region.

##  II Methodology

A normalizing flow is a density estimation model designed to learn a map
$f:\mathbb{R}^{d}\to\mathbb{R}^{d}$ between an unknown training data
distribution $X\sim p_{D}$ and a known base distribution $Z\sim p_{B}$ of the
same dimension222See Ref. Kobyzev  _et al._ (2021) for a review of modern
methods. The base distribution is typically taken to be a multidimensional
standard normal $\mathcal{N}(\mathbf{0},\mathbf{1})$, and the function $f$ is
constructed from a composition of invertible maps $f(\mathbf{x})=f_{N}\circ
f_{N-1}\circ\cdots f_{1}(\mathbf{x})$ with tractable Jacobian determinants.
The invertible structure enables sampling the unknown distribution by sampling
the base distribution, and the change of variables formula
$p_{D}(\mathbf{x})=p_{B}(f(\mathbf{x}))\left|\det\frac{\mathrm{d}f}{\mathrm{d}\mathbf{x}}\right|$
enables density estimation. Flows are trained with a log likelihood loss,
which for a composite transformation, takes the form

$\log p_{D}(\mathbf{x})=\log
p_{B}(f(\mathbf{z}))+\sum_{i=1}^{N}\log\left|\det\frac{\mathrm{d}f_{i}}{\mathrm{d}\mathbf{z}_{i-1}}\right|$

where $\mathbf{x}=\mathbf{z}_{0}$, $\mathbf{z}=\mathbf{z}_{N}$, and
$\mathbf{z}_{i}=f_{i}(\mathbf{z}_{i-1})$. The transformations $f_{i}$ are
typically drawn from a parametrized family of functions $f_{\phi}$, with
parameters $\phi$ computed by neural networks.

###  II.1 Conditional Flows & Quantile Morphing

Flow models can easily be modified to fit conditional distributions
$p(\mathbf{x}|\mathbf{y})$ by allowing the parameters of the flow
transformations $f_{i,\phi_{i}}$ to depend on the conditioning variables, i.e.
$\phi_{i}=\phi_{i}(\mathbf{x},\mathbf{y})$. A multidimensional joint density
estimation task can be decomposed into a series of one-dimensional tasks using
the probability chain rule

$p(\mathbf{x})=p(x_{1})p(x_{2}|x_{1})\cdots p(x_{d}|x_{1},\ldots,x_{d-1}),$
(1)

with each term modeled by a normalizing flow. While modern flow architectures
can perform the joint estimation, this decomposition enables the analysis of
(conditional) cumulative distribution functions (CDFs) for each dimension of
$\mathbf{x}$ and is essential to the quantile morphing technique described in
this paper.

Quantile morphing is a method to correct samples from a reference distribution
$p_{\mathrm{MC}}$ to match those of a target distribution $p_{D}$333The
subscripts MC and D are used here and throughout the paper as shorthand for
Monte Carlo and data, respectively. by applying a CDF transformation that
matches their quantiles. In 1D, the transformation $x\mapsto
y=F_{D}^{-1}(F_{\mathrm{MC}}(x))$ maps $x\sim p_{\mathrm{MC}}$ to $y\sim
p_{D}$, where $F_{\mathrm{MC}}$, $F_{D}$ are the CDFs for $p_{\mathrm{MC}}$,
$p_{D}$. This is an exact transformation from $p_{\mathrm{MC}}$ to $p_{D}$,
and guarantees optimal transport for each shifted sample.

In higher dimensions, CDFs are not uniquely defined since rotations across
dimensions are permissible, and the basic quantile morphing strategy breaks
down. Fortunately, we can reconcile this ambiguity to transform
$p_{\mathrm{MC}}$ into $p_{D}$ by breaking the problem into a series of 1D
transformations following Eq. 1 via an iterative procedure called chained
quantile morphing (CQM) Collaboration (2022); Spyromitros-Xioufis  _et al._
(2016). At each step, a dimension $x_{i}\in\mathbf{x}\sim p_{D}$ is morphed to
$y_{i}\sim p_{D}(y_{i})$ using the conditional quantile functions
$F_{\mathrm{MC}}(x_{i}|y_{1},y_{2},\ldots,y_{i-1})$ and
$F_{D}(y_{i}|y_{1},y_{2},\ldots,y_{i-1})$. Starting from the first dimension
$x_{1}$, chained quantile morphing proceeds as follows:

  1. 1.

Transform $x_{1}$ to $y_{1}=F_{D}^{-1}(F_{\mathrm{MC}}(x_{1}))$ using CDFs
$F_{\mathrm{MC}}$ and $F_{D}$.

  2. 2.

Transform $x_{2}$ to $y_{2}=F_{D}^{-1}(F_{\mathrm{MC}}(x_{2}|y_{1})|y_{1})$
using conditional CDFs $F_{\mathrm{MC}}(\cdot|y_{1})$ and $F_{D}(\cdot|y_{1})$
and the corrected value $y_{1}$ from step (1).

  3. 3.

Continue as in (2) for $i=3,\ldots,d$ with
$y_{i}=F_{D}^{-1}(F_{\mathrm{MC}}(x_{i}|\mathbf{y}_{1:i-1})|\mathbf{y}_{1:i-1})$
and the previously corrected dimensions $\mathbf{y}_{1:i-1}$.

If all (conditional) CDFs are known analytically, the CQM procedure guarantees
an exact transformation of samples $\mathbf{x}\sim p_{\mathrm{MC}}$ to
$\mathbf{y}\sim p_{D}$. In any real world application, however, the reference
and target datasets will come from complex and/or high-dimensional
distributions with intractable densities and CDFs. In these cases, the CDFs
can only be approximated, typically by training machine learning (ML)
algorithms to perform conditional quantile regression.

CQM was first introduced for LHC physics for photons within the CMS detector
in Ref. Collaboration (2022), where it was used to correct mis-modeled MC to
better match experimental data through the use of a well defined control
region within data. The authors constructed a discretized approximation of
each CDF by training boosted decision trees (BDTs) to learn a fixed array of
conditional quantiles. This was effective for the analysis but difficult to
scale efficiently due to the fixed quantiles and the large number of BDT
trainings required. In this work, we propose a streamlined approach to CQM
with normalizing flows (NFs). Flows can learn the full conditional density
(and thus CDF) at each step of CQM, enabling a more exact version of the
transformation.

###  II.2 CQM with Normalizing Flows

Figure 1: A schematic demonstrating how chained quantile morphing transforms
samples from one $k$-dimensional PDF $p_{\mathrm{MC}}(x_{1},\ldots,x_{k})$ to
another $p_{D}(x_{1},\ldots,x_{k})$. The shaded regions in blue, green, and
orange denote the conditional quantiles of variables $x_{i}$, $z_{i}$,
$y_{i}$, and the red X’s mark their values. CQM matches the conditional
quantiles of the original distribution $p_{\mathrm{MC}}$ to those of the
target distribution $p_{\mathrm{D}}$ by mapping
$x_{i}=q_{\mathrm{MC}}(\alpha_{i}|\mathbf{y}_{1:i-1})\mapsto
y_{i}=q_{\mathrm{D}}(\alpha_{i}|\mathbf{y}_{1:i-1})$, preserving the $p$-value
$\alpha_{i}$.

Figure 1 shows a schematic of the CQM correction procedure implemented with
normalizing flows. Given datasets $\\{\mathbf{x}_{j}\\}_{\mathrm{MC}}$ (“Monte
Carlo”) and $\\{\mathbf{y}_{j}\\}_{D}$ (“data”), the components
$x_{i}\in\mathbf{x}\sim p_{\mathrm{MC}}$ are iteratively corrected to match
the distributions of $y_{i}\in\mathbf{y}\sim p_{D}$ using the conditional
probability decomposition shown in Eq. 1. At step $i$, flows $f_{i,S}$ and
$f_{i,D}$ are trained to fit the conditional distributions
$p_{\mathrm{MC}}(x_{i}|\mathbf{x}_{1:i-1}^{\mathrm{corr}})$ and
$p_{D}(y_{i}|\mathbf{y}_{1:i-1})$, respectively, where
$\mathbf{x}_{1:i-1}^{\mathrm{corr}}$ are the dimensions of $\mathbf{x}$
corrected in previous steps. Individual points $x_{i,k}$ are then corrected by
the transformation

$x_{i,k}^{\mathrm{corr}}=f_{i,D}^{-1}(f_{i,\mathrm{MC}}(x_{i,k}|\mathbf{x}_{1:i-1,k}^{\mathrm{corr}})|\mathbf{x}_{1:i-1,k}^{\mathrm{corr}})$
(2)

This procedure leverages the conditional flow in lieu of the conditional
quantile function, but the quantile morphing operation is fundamentally the
same as described in Sec. II.1. A given data point $x_{i,k}\in\mathbf{x}_{k}$
will correspond to some conditional quantile of the distribution
$p_{\mathrm{MC}}(x_{i}|\mathbf{x}^{\mathrm{corr}}_{1:i-1})$, and a faithfully
trained conditional flow $f_{i,S}$ will map it to the same quantile $z_{i,k}$
of a standard normal distribution $\mathcal{N}(0,1)$. The inverse flow
$f_{i,D}$ will then map $z_{i,k}$ to the corresponding quantile of
$p_{D}(y_{i}|\mathbf{x}_{1:i-1}^{\mathrm{corr}})$. After the full chain of
corrections is applied, the samples $\\{\mathbf{x}_{j}^{\mathrm{corr}}\\}$
will follow the target distribution $p_{D}$.

###  II.3 Flow Implementation

We parameterize our flow transformations $f_{\phi}$ using piecewise rational
quadratic splines Durkan  _et al._ (2019) implemented with PyTorch Paszke  _et
al._ (2019) in the nflows package Durkan  _et al._ (2020). The splines are
defined on the interval $[-3.2,3.2]$, and all input data are min-max scaled to
the range $[-3,3]$ to capture the tails. All flows are trained with the AdamW
optimizer Kingma and Ba (2017) on a cosine-annealed learning rate schedule
Loshchilov and Hutter (2017). The spline parameters for each flow
transformation are determined from the conditioning inputs using a neural
network444When there are no conditioning inputs (i.e. the first flow of the
chain), the network is simply passed zeros..

Although we implement CQM with a sequence of distinct flows, we note that it
is possible to use an autoregressive architecture such as MADE Germain  _et
al._ (2015) to train a single flow that simultaneously learns each conditional
density of Eq. 1. While this streamlines the process, we found that it was
generally quite difficult to achieve simultaneous high-quality fits to all
inputs using this approach.

###  II.4 Comparison with Existing Strategies

Our approach to chained quantile morphing builds on the previous approach with
BDTs Collaboration (2022), and extends the scope of this effort through the
use of conditional normalizing flows. Additionally, a large variety of machine
learning-based morphing and reweighting strategies have emerged in recent
years Cranmer  _et al._ (2015); Andreassen and Nachman (2020); Andreassen  _et
al._ (2021); Nachman and Thaler (2020, 2022); Andreassen  _et al._ (2020);
Klein  _et al._ (2022); Golling  _et al._ (2023a, b); Raine  _et al._ (2023);
Sengupta  _et al._ (2023); Algren  _et al._ (2023), including several
proposals using normalizing flows Klein  _et al._ (2022); Golling  _et al._
(2023a, b); Raine  _et al._ (2023); Sengupta  _et al._ (2023); Algren  _et
al._ (2023)555For an overview of newer methods, see Ref. Golling  _et al._
(2023c). Reweighting methods aim to correct MC to match data by training a
supervised learning algorithm to differentiate the two datasets, then
reweighting the simulated dataset based on the discriminator output Hastie
_et al._ (2009); Sugiyama  _et al._ (2012). Reweighting and quantile-morphing
methods have strengths for different types of corrections, which we briefly
investigate in App. A. Broadly speaking, we find reweighting to be more
effective in correcting for theoretical physics generator mis-modeling and
quantile morphing more effective for fixing simulated detector mis-modeling.

CQM is most similar to the “flows for flows” (FFF) paradigm, proposed in Ref.
Klein  _et al._ (2022) and rephrased several times in Golling  _et al._
(2023a, b, c); Sengupta  _et al._ (2023), but differs in two key ways. First,
FFF transforms $\mathbf{x}\sim p_{\mathrm{MC}}$ to $\mathbf{y}\sim p_{D}$ all
at once using joint density estimation, rather than transforming the variables
$x_{i}$ iteratively as in CQM. Secondly, FFF operates by using the learned
density $p_{D}$ as a target for training a flow transformation that maps
$p_{\mathrm{MC}}$ directly to $p_{D}$. CQM trains separate flows for
$p_{\mathrm{MC}}$ and $p_{D}$, and transforms events $\mathbf{x}\sim
p_{\mathrm{MC}}$ via the shared Gaussian latent space. With FFF, any
unsatisfactory fit would require a re-training of the full joint density,
whereas CQM requires a single variable re-training. In cases where a very high
quality fit is needed – minimizing systematic uncertanties, for example – this
level of fine-tuned control may be advantageous.CQM additionally imposes a
notion of local transformations, as each 1D quantile map will typically shift
a variable by a relatively small distance. Other methods, such as FFF, can
only target locality by introducing an ad-hoc modification to the loss
function during training Golling  _et al._ (2023a).

##  III Experiments

Chained quantile morphing is a general-purpose technique for mapping between
any two densities, and can applied in a wide variety of contexts. In this
section, we explore its performance in two use cases: a toy example mapping
between two-dimensional datasets and a high-energy physics application using
high-level observables from simulated proton-proton collisions. In the latter
example we also explore how the transformed samples from CQM can be used in
downstream applications, namely training neural networks.

###  III.1 Toy 2D Dataset

As a simple example, we first demonstrate CQM on a pair of two-dimensional
datasets shown in Fig. 2. The top row shows the reference (half moons) and
target (checkerboard) densities, and the bottom row shows each step of the CQM
transformation. In the first step the $x$ distributions are matched, notably
thinning the density in the parts of each half moon that overlap in the $y$
direction. In the second step the $y$ distributions are matched, faithfully
reproducing the sharp boundaries and alternating conditional pattern of the
target density.

This example demonstrates the flexibility of flow-based CQM, and highlights
the advantage of using a density estimation approach that can readily capture
complex single-variable distributions and strong conditional dependencies
between inputs. Most realistic applications will have weaker correlations
between input dimensions and less exotic marginals.

Figure 2: Using CQM to morph between a pair of two-dimensional datasets (top
row). The bottom row shows the transformed distribution after morphing the
$x$-axis (left) then the $y$-axis (right).

###  III.2 LHC Olympics Datasets

We now demonstrate how CQM can be applied in high-energy physics using the LHC
Olympics (LHCO) datasets Kasieczka  _et al._ (2021). We focus on the LHCO
“R&D” Kasieczka  _et al._ (2019a) and “Black Box 1” (BB1) Kasieczka  _et al._
(2019b) datasets, and train CQM to map between them. The R&D dataset contains
one million Standard Model quantum chromodynamics (QCD) background dijet
events, and 100,000 signal dijet events featuring a heavy resonance decay
$Z^{\prime}\to XY$ ($X\to q\bar{q}$, $Y\to q\bar{q}$) with $m_{Z}=3.5$ TeV Kim
_et al._ (2020). The BB1 dataset contains a total of one million events, 834
of which are signal events with the same topology as the R&D signal but with
different resonance masses. Both datasets are generated with Pythia 8.219
Bierlich  _et al._ (2022); Sjostrand  _et al._ (2006) and Delphes 3.4.1 de
Favereau  _et al._ (2014); Mertens (2015), but with different parameter
configurations that alter the distributions of relevant high-level jet
observables.

We cluster the particles in each event into jets using the anti-$k_{T}$
algorithm Cacciari  _et al._ (2008) with radius $R=0.8$, and define jet 1 (2)
to be the heaviest (second heaviest) jet in the event. For our high-level
observables, we compute $\rho=m_{J}/p_{T,J}$ and the $N$-subjettiness ratios
$\tau_{21}$, $\tau_{32}$, and $\tau_{43}$ Thaler and Tilburg (2011) as a
measure of their 2, 3, or 4-pronged substructure.

Using only QCD background events from each dataset (i.e. removing the signal
events from BB1), we train the CQM procedure to map between them in the eight-
dimensional space constructed from $(\rho,\tau_{21},\tau_{32},\tau_{43})$ of
each jet.666Quantile morphing can be performed in either direction due to the
invertibility of normalizing flows, i.e. R&D $\to$ BB1 or BB1 $\to$ R&D We use
500,000 R&D and 500,000 BB1 events to train CQM, and evaluate its performance
with the remaining events.

Fig. 3 shows the distributions of each input variable from the test sets
before and after morphing R&D to match BB1. The morphed distributions are a
very good visual match to the target distributions, and the ratio plots
indicate good agreement across nearly the entire range of each input.
Performance begins to break down in the extreme high and low tails, but this
is to be expected due to statistical fluctuations and limited training samples
available in these sparse regions. We also show results of a standard neural
network-based reweighting scheme, implemented with a six-layer fully connected
network with ReLU activations, a Sigmoid output, and 10% dropout to prevent
overfitting. The network is trained to discriminate R&D events from BB1 events
using a binary cross entropy loss, and the score $s(\mathbf{x})$ is converted
into an event weight $w(\mathbf{x})=s(\mathbf{x})/(1-s(\mathbf{x}))$ Hastie
_et al._ (2009); Sugiyama  _et al._ (2012). As expected, the reweighting
performs on-par with CQM. Unlike CQM, however, it does not produce a set of
transformed samples that can be used for downstream tasks.

Figure 3: Distributions of $\rho,\tau_{21},\tau_{32},\tau_{43}$ (left to
right) of the leading (top row) and sub-leading (bottom row) jets in the BB1
(green), R&D (red), and morphed R&D (blue) datasets. The morphed R&D data is a
very close match to BB1, as shown by the ratio plots in the bottom panes. We
also show the results of a neural network-based reweighting of R&D data, which
yields similarly good agreement.

To better quantify the quality of the morphed distribution, we train neural
networks to distinguish BB1 samples from R&D samples before and after
corrections are applied. In Table 1 we report the area under the curve (AUC)
metrics, each averaged over an ensemble of 10 trainings. The AUCs drop to 0.51
after morphing in either direction, indicating that CQM produces samples that
are virtually indistinguishable from the target distribution. We also evaluate
the reweighting by computing weighted AUCs for the nominal training (i.e. un-
morphed samples), and find that it also succeeds.

Comparison | AUC
---|---
BB1 vs. R&D | $0.6257\pm 0.0002$
Morphed BB1 vs. R&D | $0.5093\pm 0.0004$
BB1 vs. Morphed R&D | $0.5106\pm 0.0004$
BB1 vs. Reweighted R&D | $0.5158\pm 0.0001$
Table 1: Area under the curve (AUC) metrics for neural network discriminators
trained to distinguish pairs of datasets (BB1 and R&D) before and after
applying chained quantile morphing to make one match the other. The reported
AUCs are computed from the average of 10 separate trainings. We also evaluate
re-weighting with a weighted AUC using the BB1 vs. R&D training.

###  III.3 Morphed Distributions for Classification Tasks

As demonstrated in the previous section, chained quantile morphing can
transform samples from the reference distribution into samples
indistinguishable from the target distribution. This is a powerful tool for
correcting simulation inaccuracies in high-level observables, and can find
immediate use in reducing systematic uncertainties and easing cut-based
analysis workflows. Beyond these standalone applications, however, CQM-
transformed simulation can be used in more complex downstream tasks such as
training neural networks to separate signal and background.

We demonstrate this use case using the same LHCO samples analyzed in Sec.
III.2. To emulate an LHC analysis, we treat the R&D background dataset as
“background simulation”, the R&D signal as “signal simulation”, and the BB1
background as “data”. We use CQM to construct a morphed R&D background dataset
to match the BB1 (“data”) background, and train separate neural network
classifiers to distinguish the morphed and un-morphed samples from the R&D
signal777See Sec. V for a similar study including signal contamination in the
target “data” sample.. We evaluate each classifier on the four background
samples: R&D, BB1, morphed R&D, and morphed BB1. This provides insight into
(a) how classifier performance changes when moving from “simulation” to
“data”, and (b) if training/testing on morphed samples replicates the
performance on the target datasets they are meant to match.

In Fig. 4, we plot significance improvement characteristic (SIC) curves for
each training/evaluation, where
$\mathrm{SIC}=\frac{\mathrm{TPR}}{\sqrt{\mathrm{FPR}}}$. The upper plot shows
results for trainings using the original R&D background, and colored lines
correspond evaluations on the various background datasets. The lower plot
shows the same thing, but for trainings using the morphed R&D background. For
a fixed evaluation dataset, the performance is not significantly different
between trainings, but is not necessarily expected to be. Crucially, the
performance is extremely consistent between BB1 and morphed R&D, and between
R&D and morphed BB1. In the LHC physics context, CQM would allow us to predict
a classifier’s performance on data (BB1) by morphing MC simulation (R&D) to
match data. Curiously, the classification performance is uniformly better on
the R&D/morphed BB1 samples than the BB1/morphed R&D samples. This simply
indicates that the R&D background is more separable from signal than the BB1
background 888This is because the $\tau_{21}$ is skewed lower in the BB1
sample, as shown in Fig. 3, and both jets in the $Z^{\prime}\to XY$ signal are
expected to have two prongs (i.e. low $\tau_{21}$).

Although this example is relatively simplistic (a supervised search for a
resonant signal), it demonstrates CQM’s potential for transforming imperfect
simulations into a readily usable emulation of data. This enables a reliable
estimate of an analysis’ performance on real data and relies solely on having
a suitable reference dataset, such as a signal free control region with
similar background kinematics to the signal region. As we will discuss in Sec.
V, the technique is robust against modest signal contamination in the
reference dataset and can be configured to interpolate from a kinematic
sideband into a signal region by conditioning on the relevant variables
defining the sideband (e.g. dijet mass $M_{jj}$).

Figure 4:  Significance improvement characteristic (SIC) curves for neural
networks trained to distinguish R&D background from R&D signal (top), and
morphed R&D background from R&D signal (bottom). Each curve is computed by
evaluating the classifier on a mixture of R&D signal events with background-
only BB1 (red), R&D background (purple), morphed background-only BB1 (green),
and morphed R&D background (teal).

##  IV High-Dimensional CQM using Contrastive Learning

###  IV.1 Handling large input spaces

Chained quantile morphing is ideally suited for correcting
$\mathcal{O}(10)$-dimensional input spaces constructed from high-level physics
observables. The addition of more variables complicates the construction of
the conditional probability distributions, making it harder to learn. An input
space that is $\mathcal{O}(10)$-dimensional is already very useful for many
LHC physics applications that rely on high-level observables such as precision
measurements. However, a generalized high-dimensional approach does not
naturally result from this approach. Many of the cutting-edge jet taggers
deployed at ATLAS and CMS use low-level particle- or detector-level inputs and
therefore have $\mathcal{O}(1000)$-dimensional input spaces. Applying CQM to
these inputs is infeasible due to the number of trainings required and the
fundamental ambiguity of “correcting” a variable-length unordered collection
of particle kinematics.

Despite CQM being ineffective here, resolving data/MC discrepancies in low-
level taggers is of particular interest since granular event information often
amplifies the impact of MC mis-modeling and leads to larger systematic
uncertainties. While low-level taggers perform significantly better than high-
level taggers, the large mis-modeling can almost completely remove the
improvement present. To overcome this, we rely on new approaches in deep
metric learning to learn an intermediate low-dimensional space with compelling
features that would be effective for CQM. Given such an embedding, data/MC
discrepancies could be addressed by applying CQM in the feature space, rather
than on the level of individual particles.

We propose to perform this with contrastive learning algorithms Bromley  _et
al._ (1993); Chopra  _et al._ (2005); Chen  _et al._ (2020); Dillon  _et al._
(2021, 2022), which learn through self-supervision to embed input data into a
structured feature space where different classes of inputs (e.g. light quark
jets vs. $b$ jets) are well separated. Assuming that the feature space
captures the relevant classification information from the low-level inputs, an
ML tagger trained on the resulting feature space should match the performance
of a tagger trained directly on the low-level inputs.

###  IV.2 LHCO implementation and experiments

To demonstrate this approach, we train a contrastive feature space to separate
signal and background in the LHCO R&D dataset. As low-level inputs, we use the
$(p_{T},z,\Delta\eta,\Delta\phi)$ of the 100 highest-$p_{T}$ constituents of
the leading (heaviest) jet in each event, where $z=p_{T}/p_{T,\mathrm{jet}}$,
$\Delta\eta=\eta-\eta_{\mathrm{jet}}$, and
$\Delta\phi=\phi-\phi_{\mathrm{jet}}$. We use the ParticleNet Qu and Gouskos
(2020) architecture as an embedding function into a four-dimensional feature
space, and train it on a mixture of signal and background events using the
VICReg (“Variance-Invariance-Covariance Regularization”) contrastive loss
Bardes  _et al._ (2022). VICReg learns a non-trivial representation by
rewarding functions that embed “like pairs” near to one another (invariance),
while maintaining decorrelated and suitably spread out features
(variance/covariance). In typical contrastive learning setups, training pairs
correspond to a training input and an “augmented” version of that input (e.g.
an image and a randomly rotated version of the same image). In our
implementation, we simply use pairs of signal and pairs of background events
to train the algorithm. This method effectively constructs a mostly
decorrelated feature space that captures the separation between signal and
background. The construction of a largely decorrelated feature space is
appealing for CQM, and we can use this feature space as an intermediate space
to correct data to simulation.

To assess the quality of the contrastive space, we train an MLP classifier to
distinguish embedded signal/background events and compare its performance with
that of a ParticleNet classifier trained directly on the low-level features.
We then use CQM to morph embedded R&D background events to match embedded BB1
background events and evaluate them using the MLP classifier. Finally, we
train an MLP tagger using only the high-level variables
$(\rho,\tau_{21},\tau_{32},\tau_{43})$ of the leading jet to serve as a
benchmark. The results are summarized in Table 2, and the embedding spaces are
shown in Fig. 5. The tagger trained on embeddings matches the performance of
the low-level tagger, and applying CQM in the contrastive space transforms the
R&D events to closely match BB1. By combining contrastive embedding and CQM,
we are able to retain the advantages of using low-level inputs while
simultaneously being able to correct undesired differences between two
samples. In the LHC context, this is a promising new avenue for correcting
data/MC discrepancies in low-level jet taggers and reducing the associated
systematic uncertainties.

Classifier | Sample | AUC
---|---|---
ParticleNet Tagger | R&D | 0.954
BB1 | 0.931
Embeddings (MLP) | R&D | 0.952
BB1 | 0.929
Morphed R&D | 0.928
High-Level (MLP) | R&D | 0.910
BB1 | 0.872
Morphed R&D | 0.873
Table 2: A table showing the performance of the three taggers described in the
text, evaluated by the AUC for separating R&D signal from the indicated
background sample. ParticleNet (top) is trained on low-level features, where
as the other two are trained on contrastive embeddings (middle) and high-level
jet features (bottom).

Figure 5:  (Left) The resulting latent space constructed from self-supervised
VICReg training using the R&D background sample and the $Z^{\prime}\to XY$
signal. The red shows embeddings of the background-only BB1 sample appears,
which disagree slightly with the R&D backgrounds. (Right) The same plot, but
after applying CQM to transform the R&D background to match the BB1 background
in the latent space.

##  V Signal Contamination & Sideband Studies

The studies presented thus far morph an inclusive “MC” background sample to an
inclusive “data” sample with the assumption that the underlying physics in
both samples are the same. In most realistic LHC analyses, however, MC samples
are used to estimate backgrounds in regions of phase space where the data is
expected to contain (in the case of measurements) or may contain (in the case
of searches) additional physics unaccounted for by the simulations. It is thus
unrealistic to simply morph between two inclusive datasets. Instead, we would
like to use a signal-free control region to learn the morphing functions, and
then interpolate or extrapolate these functions into the signal region. In
this section, we present two studies addressing this question.

First, we consider the case of interpolating the morphing function trained in
a control region. As in Refs. Golling  _et al._ (2023b, a, c) and in keeping
with the LHCO dataset, we consider the case of a resonant signal decaying to a
pair of jets. We take the resonant mass to be $M_{\mathrm{res}}=3.5$ TeV,
define our signal region (SR) as $M_{jj}\in[3.3,3.7]$ TeV, and train CQM in
the $M_{jj}$ sideband (SB). We again treat the R&D dataset as “MC” and signal-
free BB1 as “data”, and transform the variables
$(\rho,\tau_{21},\tau_{32},\tau_{43})$ of the leading jet. All flows are
conditioned on $M_{jj}$ to allow interpolation into the signal region. In
Table 3 (top row), we show AUCs for a classifier trained to distinguish
signal-free BB1 from morphed R&D in the SB and SR. The AUCs indicate that they
are indistinguishable in both the sideband and signal region, demonstrating
that CQM remains effective when interpolating into a region of phase space
unseen during training. In this case, the SR was defined by a cut on a single
variable, $M_{jj}$, but this could easily be generalized to SRs defined by
multiple observables.

As a second test, we consider the case where the signal is small relative to
the background. If the signal is sufficiently small, it contributes so
negligibly to the single variable “data” density learned by CQM that it is
unnecessary to define cuts that remove it. To test this, we train using the
R&D background sample as “MC” and the full BB1 dataset as “data” – i.e.
without signal events artificially removed – with a signal contamination
fraction of 0.08%. Table 3 (bottom) indicates that the morphed R&D sample is
still essentially indistinguishable from the uncontaminated BB1 background in
this case, showing that a small degree of signal contamination has little to
no impact on the quality of the morphed samples. The impact of small
contamination is particularly robust for CQM when no single variable dominates
the discrimination and the discrimination is spread over a variety of
correlated variables. Since CQM is applied one variable at a time, sensitivity
to “signal regions” is particularly reduced, making CQM robust in the regions
of high separation between signal and background.

Training | Selection | AUC
---|---|---
Sideband | SB | $0.514\pm 0.001$
SR | $0.510\pm 0.004$
Contaminated | Inclusive | $0.514\pm 0.001$
Table 3: Evaluating the quality of the morphed R&D $\to$ BB1 samples when CQM
is configured to interpolate into a blinded signal region defined by an
$M_{jj}$ (top), or trained using a slightly contaminated “data” sample
(bottom). As before, the AUC is computed using 10 neural network trainings to
distinguish BB1 from morphed R&D.

##  VI Conclusion

In this paper, we have presented a normalizing flow-based implementation of
the “chained quantile morphing” (CQM) technique for correcting Monte Carlo
simulations to better match experimental data. CQM matches the performance of
reweighting with neural networks and comes with the added benefit of producing
a new set of corrected samples rather than simply event weights. Moreover,
quantile mapping is a fundamentally different approach to reweighting and can
be particularly effective when correcting for detector mis-modeling.

Our approach is unique in that the chained structure makes the implementation
and interpretation of the results markedly simple and robust. The iterative,
conditional density estimation – as opposed to a simultaneous joint estimation
– allows intervention at each stage of the morphing process, where the flow
architectures can be modified or re-trained to maximize performance. This
means we can ensure a high-quality fit to each variable without re-training
the entire morphing process for every small modification to the flow. The
structure of CQM also ensures some degree of locality in the morphing
transformations since each variable is transformed according to 1D conditional
CDFs. This is the optimal transport map for 1D data. When chained together for
an $N$-dimensional problem, it is intuitive that each variable is not moved
“too far” from its initial value. This stands in contrast to the joint density
approaches where the training ensures that the overall base density is
transformed to match the target but does not guarantee small overall
corrections per variable.

In the emerging landscape of morphing and reweighting strategies, most
applications have been focused on background estimation for resonant anomaly
searches. While this is a worthwhile application, we propose broadening the
discussion to consider how CQM – and MC correction strategies in general – can
impact the LHC physics program. As seen in this study and experimentally in
Ref. Collaboration (2022), CQM is a promising tool for reducing systematic
uncertainties by correcting mis-modeled simulations. This is useful for a wide
range of physics analyses using $\mathcal{O}(10)$ high-level inputs,
especially those that rely on neural networks or boosted decision trees to
build a classifier. We have also demonstrated CQM’s potential for problems
with very high-dimensional inputs (e.g. jet taggers), where contrastive
learning allows us to compress and correct the relevant information in a low-
dimensional feature space.

As our machine learning tools become more sophisticated, they will continue to
expose and amplify the flaws in our particle physics simulations. The LHC
physics program already relies heavily on ML tools, and we will inevitably
come to a point where the uncertainties associated with a simulation-trained
model significantly limit the sensitivity of a search or the precision of a
measurement. CQM is not a perfect solution to this problem, but it represents
a promising step towards the ultimate goal of robust and unbiased ML. Moving
ahead, we hope to apply a version of CQM to real LHC data and look forward to
the ongoing development and refinement of morphing strategies.

###### Acknowledgements.

We thank Gregor Kasieczka for helpful and constructive feedback on the
manuscript. PH, PM, and SR are supported in part by the Institute for
Artificial Intelligence and Fundamental Interactions (IAIFI) through NSF Grant
No. PHY-2019786, and the NSF Institute for AI Accelerated Algorithms for Data-
Driven Discovery (NSF Grant #PHY-2117997). Additional support comes from the
FAIR Data program of the U.S. DOE, Office of Science, Advanced Scientific
Computing Research, under contract number DE-AC02-06CH11357, and a DOE early
career award.

##  Appendix A The limits of morphing & reweighting

Both chained quantile morphing and reweighting are subject to different
intrinsic limitations, which must inform the choice of which technique to use
in a real-world application. Conceptually the style of applications are
different between the two algorithms. Morphing corresponds to a shift of the
simulated distribution to match the data. This corresponds to a detector mis-
modeling where the detector response or resolution are off and the
distribution requires a shifting and scaling to match the observed resolution.
A reweighting is an adjustment of the simulation to emphasize/de-emphaise
regions of the generated parameter space. This corresponds to correction to
the simulation to match the intrinsic generated parameter to reality. This
intuitively corresponds to generation mis-modeling such as a missing higher
order correction or shower parameter. As a result, with each there are both
advantages and disadvantages.

We have constructed two toy examples of situations where one approach may be
preferred over the other:

###  A.1 The limitations of reweighting

One of the major strengths of a morphing scheme is that it is capable of
correcting a distribution into a range not well-covered by the original
simulated sample. In our toy example, we consider a 5-dimensional dataset
where the “data” has a strong correlation between two of the dimensions that
is missed in the “MC” sample. In particular, we define the distributions

$v_{i}^{(MC)}=\mathcal{N}(0,1)\,\,\forall\,\,i$ (3) $\displaystyle
v_{i}^{(data)}$ $\displaystyle=\mathcal{N}(0,1)\,\,\forall\,\,i\neq 2$ (4)
$\displaystyle v_{2}^{(data)}$ $\displaystyle=\mathcal{N}(0,1)+5v_{3}$ (5)

Note that this creates a situation in which the ”MC” sample is not completely
overlapping with the “data” sample.

We then corrected the MC distribution onto the data distribution with both CQR
and a simple reweighting scheme. Kernel density estimate contours of the
resulting distribution are shown in Figure 6, where it is clear that the
reweighting scheme fails to reproduce the data distribution where it fails to
overlap the MC distribution. This would likely be improved with a larger
dataset where the reweighting has more examples in the MC tails to reweight
into the data distribution. In general, we find the CQR scheme to be more
robust than reweighting, particularly in matching parts of the distribution
that are not well-covered by the source distribution.

Figure 6: Kernel density estimate contours for ”MC” and ”data” distributions
for our toy example, together with the CQR-corrected distribution (top) and
the reweighting-corrected distribution (bottom).

###  A.2 The limitations of morphing

On the other hand, one of the major strengths of reweighting is that it
preserves by construction all of the complicated relationships between the
variables in any given event. In situations where the variables being
considered are related by some invariant or conserved quantity reweighting
guarantees that that relationship is preserved in the transformed
distribution. Morphing makes no such guarantee, and only preserves this
relationship insofar as it is able to learn this relationship in the training.

Figure 7: Reconstructed invariant mass distributions for ”MC” and ”data”
distributions for our toy example, together with the CQR-corrected
distribution (top) and the reweighting-corrected distribution (bottom).

In order to demonstrate this effect we consider a simplified toy model of the
decay of a narrow resonance (suggestively placed at 125 GeV). In this dataset
we simulate the decay of this resonance by first pulling an invariant mass
from a normal distribution $\mathcal{N}(125,5)$. One decay daughter is
generated by randomly generating an momentum and angle in the 2D plane, and
the second daughter is generated by randomly generating a $\delta\phi$ and
then fixing the momentum to perfectly conserve the generated invariant mass.
This creates a complex and perfectly fixed relationship between the momentum
of the second particle and the momentum of the first particle, the
$\delta\phi$ between them, and the generated invariant mass. In the “MC”
sample the momentum distribution and $\delta\phi$ distribution are perturbed
by $\approx 5\%$ with respect to the “data” sample, but the invariant masses
are drawn from exactly the same distribution. We then blind the generated
invariant mass and use both reweighting and CQR to correct the “MC”
distributions onto the “data” ones. In order to evaluate the performance of
the two schemes at preserving the invariant mass we reconstructed it from the
two decay daughters and show the resulting distributions in Figure 7. The
reweighting scheme perfectly preserves the invariant mass, while the CQR isn’t
able to perfectly learn this relationship and smears out the invariant mass
distribution. This could likely be improved with more data or more fine-tuning
of the trainings in order to help the CQR model learn the complex relationship
between the variables. An alternative approach would be to reconstruct the
invariant mass distribution in both the data and MC samples first, and then
perform a morphing of this distribution from MC to data.

## References

  * Qu and Gouskos (2020) Huilin Qu and Loukas Gouskos, “ParticleNet: Jet Tagging via Particle Clouds,” Phys. Rev. D 101, 056019 (2020), arXiv:1902.08570 [hep-ph] .
  * CMS (2023) _Performance of heavy-flavour jet identification in boosted topologies in proton-proton collisions at $\sqrt{s}=13~{}\mathrm{TeV}$_, Tech. Rep. (CERN, Geneva, 2023).
  * Dreyer and Qu (2021) Frédéric A. Dreyer and Huilin Qu, “Jet tagging in the Lund plane with graph networks,” JHEP 03, 052 (2021), arXiv:2012.08526 [hep-ph] .
  * ATL (2023) “Tagging boosted $W$ bosons applying machine learning to the Lund Jet Plane,” (2023).
  * Rothman (2021) Simon Rothman (CMS),  _Calibrating Electrons and Photons in the CMS ECAL using Graph Neural Networks_, Tech. Rep. (CERN, Geneva, 2021).
  * Spyromitros-Xioufis  _et al._ (2016) Eleftherios Spyromitros-Xioufis, Grigorios Tsoumakas, William Groves, and Ioannis Vlahavas, “Multi-target regression via input space expansion: treating targets as inputs,” Machine Learning 104, 55–98 (2016).
  * Collaboration (2022) CMS Collaboration, “Measurement of the higgs boson inclusive and differential fiducial production cross sections in the diphoton decay channel with pp collisions at $\sqrt{s}$ = 13 tev,” (2022), arXiv:2208.12279 [hep-ex] .
  * Kobyzev  _et al._ (2021) Ivan Kobyzev, Simon J.D. Prince, and Marcus A. Brubaker, “Normalizing flows: An introduction and review of current methods,” IEEE Transactions on Pattern Analysis and Machine Intelligence 43, 3964–3979 (2021).
  * (9) E. G. Tabak and Cristina V. Turner, “A family of nonparametric density estimation algorithms,” Communications on Pure and Applied Mathematics 66, 145–164, https://onlinelibrary.wiley.com/doi/pdf/10.1002/cpa.21423 .
  * Papamakarios  _et al._ (2021) George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan, “Normalizing flows for probabilistic modeling and inference,” (2021), arXiv:1912.02762 [stat.ML] .
  * Cranmer  _et al._ (2015) Kyle Cranmer, Juan Pavez, and Gilles Louppe, “Approximating Likelihood Ratios with Calibrated Discriminative Classifiers,” (2015), arXiv:1506.02169 [stat.AP] .
  * Andreassen and Nachman (2020) Anders Andreassen and Benjamin Nachman, “Neural networks for full phase-space reweighting and parameter tuning,” Phys. Rev. D 101, 091901 (2020).
  * Andreassen  _et al._ (2021) Anders Andreassen, Shih-Chieh Hsu, Benjamin Nachman, Natchanon Suaysom, and Adi Suresh, “Parameter estimation using neural networks in the presence of detector effects,” Phys. Rev. D 103, 036001 (2021).
  * Nachman and Thaler (2020) Benjamin Nachman and Jesse Thaler, “Neural resampler for monte carlo reweighting with preserved uncertainties,” Phys. Rev. D 102, 076004 (2020).
  * Nachman and Thaler (2022) Benjamin Nachman and Jesse Thaler, “Neural conditional reweighting,” Phys. Rev. D 105, 076015 (2022).
  * Andreassen  _et al._ (2020) Anders Andreassen, Benjamin Nachman, and David Shih, “Simulation Assisted Likelihood-free Anomaly Detection,” Phys. Rev. D 101, 095004 (2020), arXiv:2001.05001 [hep-ph] .
  * Klein  _et al._ (2022) Samuel Klein, John Andrew Raine, and Tobias Golling, “Flows for flows: Training normalizing flows between arbitrary distributions with maximum likelihood estimation,” (2022), arXiv:2211.02487 [cs.LG] .
  * Golling  _et al._ (2023a) Tobias Golling, Samuel Klein, Radha Mastandrea, Benjamin Nachman, and John Andrew Raine, “Flows for Flows: Morphing one Dataset into another with Maximum Likelihood Estimation,” (2023a), arXiv:2309.06472 [hep-ph] .
  * Golling  _et al._ (2023b) Tobias Golling, Samuel Klein, Radha Mastandrea, and Benjamin Nachman, “Flow-enhanced transportation for anomaly detection,” Phys. Rev. D 107, 096025 (2023b), arXiv:2212.11285 [hep-ph] .
  * Raine  _et al._ (2023) John Andrew Raine, Samuel Klein, Debajyoti Sengupta, and Tobias Golling, “CURTAINs for your sliding window: Constructing unobserved regions by transforming adjacent intervals,” Front. Big Data 6, 899345 (2023), arXiv:2203.09470 [hep-ph] .
  * Sengupta  _et al._ (2023) Debajyoti Sengupta, Samuel Klein, John Andrew Raine, and Tobias Golling, “CURTAINs Flows For Flows: Constructing Unobserved Regions with Maximum Likelihood Estimation,” (2023), arXiv:2305.04646 [hep-ph] .
  * Algren  _et al._ (2023) Malte Algren, Tobias Golling, Manuel Guth, Chris Pollard, and John Andrew Raine, “Flow Away your Differences: Conditional Normalizing Flows as an Improvement to Reweighting,” (2023), arXiv:2304.14963 [hep-ph] .
  * Durkan  _et al._ (2019) Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios, “Neural spline flows,” (2019), arXiv:1906.04032 [stat.ML] .
  * Paszke  _et al._ (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala, “Pytorch: An imperative style, high-performance deep learning library,” (2019), arXiv:1912.01703 [cs.LG] .
  * Durkan  _et al._ (2020) Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios, “nflows: normalizing flows in PyTorch,”  (2020).
  * Kingma and Ba (2017) Diederik P. Kingma and Jimmy Ba, “Adam: A method for stochastic optimization,” (2017), arXiv:1412.6980 [cs.LG] .
  * Loshchilov and Hutter (2017) Ilya Loshchilov and Frank Hutter, “Sgdr: Stochastic gradient descent with warm restarts,” (2017), arXiv:1608.03983 [cs.LG] .
  * Germain  _et al._ (2015) Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle, “Made: Masked autoencoder for distribution estimation,” (2015), arXiv:1502.03509 [cs.LG] .
  * Golling  _et al._ (2023c) Tobias Golling, Gregor Kasieczka, Claudius Krause, Radha Mastandrea, Benjamin Nachman, John Andrew Raine, Debajyoti Sengupta, David Shih, and Manuel Sommerhalder, “The Interplay of Machine Learning–based Resonant Anomaly Detection Methods,” (2023c), arXiv:2307.11157 [hep-ph] .
  * Hastie  _et al._ (2009) T. Hastie, R. Tibshirani, and J.H. Friedman,  _The Elements of Statistical Learning: Data Mining, Inference, and Prediction_, Springer series in statistics (Springer, 2009).
  * Sugiyama  _et al._ (2012) Masashi Sugiyama, Taiji Suzuki, and Takafumi Kanamori,  _Density Ratio Estimation in Machine Learning_ (Cambridge University Press, 2012).
  * Kasieczka  _et al._ (2021) Gregor Kasieczka _et al._ , “The LHC Olympics 2020 a community challenge for anomaly detection in high energy physics,” Rept. Prog. Phys. 84, 124201 (2021), arXiv:2101.08320 [hep-ph] .
  * Kasieczka  _et al._ (2019a) Gregor Kasieczka, Ben Nachman, and David Shih, “R&D Dataset for LHC Olympics 2020 Anomaly Detection Challenge,”  (2019a).
  * Kasieczka  _et al._ (2019b) Gregor Kasieczka, Benjamin Nachman, and David Shih, “Official Datasets for LHC Olympics 2020 Anomaly Detection Challenge,”  (2019b).
  * Kim  _et al._ (2020) Jeong Han Kim, Kyoungchul Kong, Benjamin Nachman, and Daniel Whiteson, “The motivation and status of two-body resonance decays after the LHC Run 2 and beyond,” JHEP 04, 030 (2020), arXiv:1907.06659 [hep-ph] .
  * Bierlich  _et al._ (2022) Christian Bierlich _et al._ , “A comprehensive guide to the physics and usage of PYTHIA 8.3,”  (2022), 10.21468/SciPostPhysCodeb.8, arXiv:2203.11601 [hep-ph] .
  * Sjostrand  _et al._ (2006) Torbjorn Sjostrand, Stephen Mrenna, and Peter Z. Skands, “PYTHIA 6.4 Physics and Manual,” JHEP 05, 026 (2006), arXiv:hep-ph/0603175 .
  * de Favereau  _et al._ (2014) J. de Favereau, C. Delaere, P. Demin, A. Giammanco, V. Lemaître, A. Mertens, and M. Selvaggi (DELPHES 3), “DELPHES 3, A modular framework for fast simulation of a generic collider experiment,” JHEP 02, 057 (2014), arXiv:1307.6346 [hep-ex] .
  * Mertens (2015) Alexandre Mertens, “New features in Delphes 3,” J. Phys. Conf. Ser. 608, 012045 (2015).
  * Cacciari  _et al._ (2008) Matteo Cacciari, Gavin P. Salam, and Gregory Soyez, “The anti-$k_{t}$ jet clustering algorithm,” JHEP 04, 063 (2008), arXiv:0802.1189 [hep-ph] .
  * Thaler and Tilburg (2011) Jesse Thaler and Ken Van Tilburg, “Identifying boosted objects with n-subjettiness,” Journal of High Energy Physics 2011 (2011), 10.1007/jhep03(2011)015.
  * Bromley  _et al._ (1993) Jane Bromley, Isabelle Guyon, Yann LeCun, Eduard Säckinger, and Roopak Shah, “Signature verification using a ”siamese” time delay neural network,” in  _Proceedings of the 6th International Conference on Neural Information Processing Systems_ , NIPS’93 (Morgan Kaufmann Publishers Inc., San Francisco, CA, USA, 1993) p. 737–744.
  * Chopra  _et al._ (2005) S. Chopra, R. Hadsell, and Y. LeCun, “Learning a similarity metric discriminatively, with application to face verification,” in  _2005 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR’05)_, Vol. 1 (2005) pp. 539–546 vol. 1.
  * Chen  _et al._ (2020) Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton, “A simple framework for contrastive learning of visual representations,” (2020), arXiv:2002.05709 [cs.LG] .
  * Dillon  _et al._ (2021) Barry M. Dillon, Gregor Kasieczka, Hans Olischlager, Tilman Plehn, Peter Sorrenson, and Lorenz Vogel, “Symmetries, safety, and self-supervision,” (2021), arXiv:2108.04253 [hep-ph] .
  * Dillon  _et al._ (2022) Barry M. Dillon, Radha Mastandrea, and Benjamin Nachman, “Self-supervised anomaly detection for new physics,” Physical Review D 106 (2022), 10.1103/physrevd.106.056005.
  * Bardes  _et al._ (2022) Adrien Bardes, Jean Ponce, and Yann LeCun, “Vicreg: Variance-invariance-covariance regularization for self-supervised learning,” (2022), arXiv:2105.04906 [cs.CV] .


# On the Exactness of an Energy-efficient Train Control model based on Convex
Optimization

Shaofeng Lu
Shien-ming Wu School of Intelligent Engineering
South China University of Technology
Guangzhou, China 511442
<EMAIL_ADDRESS>
&Minling Feng
Shien-ming Wu School of Intelligent Engineering
South China University of Technology
Guangzhou, China 511442
<EMAIL_ADDRESS>
&Kunpeng Wu
Shien-ming Wu School of Intelligent Engineering
South China University of Technology
Guangzhou, China 511442
<EMAIL_ADDRESS>
Corresponding author, https://lushaofeng.github.io

###### Abstract

In this paper, we demonstrate the exactness proof for the energy-efficient
train control (EETC) model based on convex optimization. The proof of
exactness shows that the convex optimization model will share the same
optimization results with the initial model on which the convex relaxations
are conducted. We first show how the relaxation on the initial non-convex
model is conducted and provide analysis to show that the relaxations are
convex constraints and the relaxed model is thus a convex model. Subsequently,
we prove that the relaxed convex model will always achieve its optimal
solution on the initial equality constraints and the optimal solution achieved
by convex optimization will be the same as the one obtained by the initial
non-convex model and the relaxations applied are exact. A numerical
verification has been conducted based on a typical urban rail system with a
steep gradient. The results of this paper shed lights on further applications
of convex optimization on energy-efficient train control and relevant areas
related to operation and control of low-carbon transportation systems.

##  1 Introduction

Energy-efficient train control (EETC) tries to locate the most energy-
efficient train speed trajectory resulted from the optimal control strategies
of the train dynamic system. The past studies in the field generally adopted
two main types of methods: indirect method based on the optimal theory and
direct methods based on different optimization techniques.

On the one hand, indirect method obtained its name from the way of locating
the optimal solution by using the Pontryagin’s Maximum Principle (PMP).
According to PMP, the optimal solution’s necessary condition can be obtained,
upon which the solution can be evaluated using the necessary condition. The
optimal solution can be obtained by indirectly solving the co-state variable’s
differential equations with the help of the necessary conditions. On the other
hand, direct methods try to set up mathematical models and apply computer
algorithms to locate the optimal solution. Our proposed EETC model with its
solution approach falls into the direct method and in general, we are
proposing a mathematical model solved by efficient algorithms to quickly
locate the optimal solutions. These solutions are generally represented by the
train speed and its corresponding train control strategies, i.e. acceleration,
braking and cruising etc. We refer readers to numerous influential academic
papers on both indirect methods [1, 2, 3, 4, 5] and direct methods[6, 7, 8, 9,
10].

Based on the previous research outcomes of [11], in this paper, we will
presents two direct EETC optimization models. The first model, referred to as
“Model A” is an non-convex model and the second mode is the relaxed “Model B”.
We focus on the topic how the two optimization models can share same
optimization results after the relaxation has been conducted from one to
another. We first present two models, offer some analysis on the convexity of
Model B in Section 2, then provide the exactness proof of relaxations in
Section 3 and demonstrate a numerical verification using a typical urban rail
case with a steep downhill gradient in Section 4. A conclusion is drawn in
Section 5.

##  2 Two EETC Models

###  2.1 The non-convex EETC model

A similar introduction of this model can be found in [11], for the sake of
coherence, we still include the key modeling details in this section. The
journey is divided into $N$ segments. For the sake of simplicity, the length
of each segment is equal and assumed to be $\Delta d$. The total length of
journey is denoted by $D$. We have

$D=\sum_{i=1}^{N}\Delta d_{i}$ (1)

where $i$ is the index number corresponding to each candidate speed point
along the journey.

The whole journey includes $N+1$ speed variable $v_{i}$ corresponding to $N$
segments. $v_{0}$ and $v_{N}$ are initial and final candidate speed. Each
speed $v_{i}$ is constrained by the speed limit $V_{i,lim}$ as presented by:

$0<v_{i}\leq V_{i,lim}$ (2)

where $V_{i,lim}$ is a parameter setting the speed limit for each $v_{i}$. In
the proof demonstrated later, we make a special case for $v_{0}$ and assume
$v_{0}=0$ to ensure the train departs from zero speed.

The sum of elapsed time in each segment equals to the total journey time $T$
as shown in (3).

$\sum_{i=1}^{N}\frac{\Delta d}{v_{i}}=T.$ (3)

The drag force can be calculated by the Davis equation in (4):

$f_{i,d}=A+Bv_{i}+Cv_{i}^{2}$ (4)

where the parameters $A$, $B$ and $C$ are the Davis coefficients.

The principle of energy conservation during train running can be presented by
(2.1):

$\displaystyle F_{i}\Delta d$
$\displaystyle=0.5M(v_{i}^{2}-v_{i-1}^{2})+f_{i,d}\Delta d+Mg\Delta H_{i}$
$\displaystyle=0.5M(v_{i}^{2}-v_{i-1}^{2})+(A+Bv_{i}+Cv_{i}^{2})\Delta
d+Mg\Delta H_{i}$ (5)

where $F_{i}$ represents the tractive or braking effort imposed the train
based on the principle of conservation of energy, $M$ is the total mass of the
train which can be varied to consider the rotary mass, and $g$ is acceleration
rate due to gravity. $\Delta H_{i}$ is the altitude difference between the
current distance $d_{i}$ and previous position $d_{i-1}$ and $\Delta H_{i}$ is
positive when it is going uphill.

The inequality constraints are imposed to ensure that the train effort and
power does not exceed its boundaries:

$-F_{max}\leq F_{i}\leq F_{max}$ (6)

$\frac{-P_{b,max}}{v_{i}}\leq F_{i}\leq\frac{P_{t,max}}{v_{i}}$ (7)

where $P_{b,max}$ is the maximum braking power and $P_{t,max}$ is the maximum
traction power. (6) is to denote the force boundary condition and (7) is for
the power boundary.

The electrical energy consumed during traction and generated during the
regenerative braking can be modeled by:

$\displaystyle E_{i}$ $\displaystyle\geq F_{i}\Delta d/\eta_{t}$ (8a)
$\displaystyle E_{i}$ $\displaystyle\geq F_{i}\Delta d\eta_{b}$ (8b)

where $1\geq\eta_{t}>0$ and $1\geq\eta_{b}>0$ are the motor efficiency during
traction and braking respectively. When $F_{i}$ is positive (8b) is relaxed
and otherwise (8a) is relaxed. $\eta_{b}$ can be set to sufficiently small
number so that no regenerative braking energy will be considered.

The total net energy is defined as the total positive tractive energy plus the
total negative regenerative braking energy.

The non-convex EETC model can be presented by:

$\displaystyle\mathrm{minimize}\leavevmode\nobreak\ \leavevmode\nobreak\
f_{\rm obj}(E)=\sum_{i=1}^{N}E_{i}$ (9)
$\displaystyle\mathrm{Subject\leavevmode\nobreak\ to:\leavevmode\nobreak\
}\eqref{for:jou_length_constr}-\eqref{for:elec_energy}$

Since the objective function is to minimize, $E_{i}$ will always approach
$F_{i}\Delta d\eta_{t}$ or $F_{i}\Delta d\eta_{b}$ as shown in (8).

###  2.2 The relaxed EETC model

In the relaxed convex model, two new sets of variables are introduced:
$\alpha_{i}$ and $\beta_{i}$.

We add two constraints defined in (10) and (11).

$\alpha_{i}\geq 1/v_{i}\mathrm{\leavevmode\nobreak\ \leavevmode\nobreak\
for\leavevmode\nobreak\ \leavevmode\nobreak\ i=1,2,\cdots,N}$ (10)

$\beta_{i}\geq v_{i}^{2}\mathrm{\leavevmode\nobreak\ \leavevmode\nobreak\
for\leavevmode\nobreak\ \leavevmode\nobreak\ i=1,2,\cdots,N-1}$ (11)

The constraint (2) is transformed into (12):

$\beta_{i}\leq V_{i,lim}^{2}.$ (12)

The constraint (3) is transformed into (13):

$\sum_{i=1}^{N}\Delta d\alpha_{i}=T.$ (13)

The constraint (4) is transformed into (14):

$f_{i,d}=A+Bv_{i}+C\beta_{i}.$ (14)

The constraint (2.1) is transformed into (2.2):

$\displaystyle F_{i}\Delta d$
$\displaystyle=0.5M(\beta_{i}-\beta_{i-1})+f_{i,d}\Delta d+Mg\Delta H_{i}$
$\displaystyle=0.5M(\beta_{i}-\beta_{i-1})+(A+Bv_{i}+C\beta_{i})\Delta
d+Mg\Delta H_{i}.$ (15)

The power boundary constraint (7) is transformed into (16):

$-P_{b,max}\alpha_{i}\leq F_{i}\leq P_{t,max}\alpha_{i}.$ (16)

The relaxed EETC model can be presented by (17):

$\displaystyle\mathrm{minimize}\leavevmode\nobreak\ \leavevmode\nobreak\
f_{\rm obj}(E)=\sum_{i=1}^{N}E_{i}$ (17)
$\displaystyle\mathrm{Subject\leavevmode\nobreak\ to:\leavevmode\nobreak\
}\eqref{for:jou_length_constr},\eqref{for:force_boundary},\eqref{for:elec_energy},\eqref{for:alpha_i}-\eqref{for:power_boundary_cvx}.$

The relaxed EETC model is built with two relaxation constraints defined in
(10) and (11). In Subsection 2.3, we demonstrate that these two constraints
are convex and the new EETC defined in (17) is a convex optimization model. In
Section 3, we prove that the two relaxation constraints in (10) and (11) are
exact. In other words, the optimal solution will always be attained on the
equality of the constraints.

###  2.3 Analysis of convexity of two relaxed constraints

The constraint (10) can be easily transformed into a Second-order Conic
Program (SOCP) given that both $\alpha_{i}\geq 0$ and $v_{i}\geq 0$.

We conduct the following transformation on (10) as presented in (18):

$\displaystyle 1$ $\displaystyle\leq\alpha_{i}v_{i}$ (18a) $\displaystyle
2^{2}+(\alpha_{i}-v_{i})^{2}$
$\displaystyle\leq\alpha_{i}^{2}+v_{i}^{2}+2\alpha_{i}v_{i}$ (18b)
$\displaystyle\sqrt{2^{2}+(\alpha_{i}-v_{i})^{2}}$
$\displaystyle\leq\alpha_{i}+v_{i}.$ (18c)

Define new variables $x_{i}$, $y_{i}$ and $z_{i}$ in (19):

$\displaystyle x_{i}$ $\displaystyle=2$ (19a) $\displaystyle y_{i}$
$\displaystyle=\alpha_{i}-v_{i}$ (19b) $\displaystyle z_{i}$
$\displaystyle=\alpha_{i}+v_{i}.$ (19c)

This transforms the program (18) into a standard SOCP as shown in (20):

$\displaystyle\sqrt{x_{i}^{2}+y_{i}^{2}}\leq z_{i}.$ (20)

The definition in (19) represents the intersection of a second-order cone with
affine sets and this does not change the convexity of the constraint in (18).
It thus verifies that constraint (10) is a convex SOCP constraint given that
$v_{i}>0$ and $\alpha_{i}>0$.

On the other hand, (11) is an inequality involving quadratic term $v_{i}^{2}$
and it is a convex constraint which can be verified by the first-order
condition of convex function [12, Page 83]. For the sake of simplicity, we
consider the following function to represent $v_{i}^{2}\leq\beta_{i}$:

$f(\mathbf{x})=x_{1}^{2}-x_{2}\leavevmode\nobreak\ \leavevmode\nobreak\
\leavevmode\nobreak\ x_{1}\geq 0\leavevmode\nobreak\ \leavevmode\nobreak\
\mathrm{and}\leavevmode\nobreak\ \leavevmode\nobreak\ x_{2}\geq 0$ (21)

where $\mathbf{x}=[x_{1},x_{2}]^{T}$ is to represent the decision vector
$[v_{i},\beta_{i}]^{T}$.

The first-order condition of convexity states that if function $f$ is
differentiable, then $f$ is convex if and only if the feasible domain of $f$
is convex and

$f(y)\geq f(x)+\nabla f(x)^{T}(y-x).$ (22)

The feasible domain of function $f(\mathbf{x})$ is convex which can be proved
by the convex set definition. We conduct the following deduction on (21) based
on the first-order conditions for convexity:

$\displaystyle f_{i}(\mathbf{y_{i}})-[f_{i}(\mathbf{x_{i}})+\nabla
f(x)^{T}(\mathbf{y_{i}}-\mathbf{x_{i}})]$
$\displaystyle=y_{1}^{2}-y_{2}-[x_{1}^{2}-x_{2}+2x_{1}(y_{1}-x_{1})-(y_{2}-x_{2})]$
$\displaystyle=y_{1}^{2}-y_{2}-(-y_{2}-x_{1}^{2}+2x_{1}y_{1})$
$\displaystyle=(y_{1}-x_{1})^{2}\geq 0$ (23)

This shows that the constraint $v_{i}^{2}\leq\beta_{i}$ is a convex
constraint. It can be seen that all constraints in Model B is either affine
equality or convex inequality. The objective function (17) is also an affine
function. Optimization problem based on Model B is a convex optimization
problem.

##  3 Proof of Exactness of Relaxations

We use Model A to refer to the model defined by (9) and Model B to refer to
the model defined by (17).

######  Theorem 1.

The optimal solution of Model A is also the optimal solution of Model B and
thus the relaxation constraints (10) and (11) in Model B are exact.

Assume the optimal solution for the model defined by Model A is
$\mathcal{S}_{1}^{*}$ and the optimal solution for Model B is
$\mathcal{S}_{2}^{*}$.

Since Model B is a relaxed model of Model A,
$f_{\mathrm{obj}}(\mathcal{S}_{2}^{*})\leq
f_{\mathrm{obj}}(\mathcal{S}_{1}^{*})$. On the other hand, if the optimal
solution of Model B is always a feasible solution of Model A, it can be
deducted that $f_{\mathrm{obj}}(\mathcal{S}_{1}^{*})\leq
f_{\mathrm{obj}}(\mathcal{S}_{2}^{*})$. This assumption is valid as long as
the optimal solution can always be attained on the equality of (10) and (11).
As a result, when the optimal solution is achieved on the equality of (10) and
(11), we have
$f_{\mathrm{obj}}(\mathcal{S}_{1}^{*})=f_{\mathrm{obj}}(\mathcal{S}_{2}^{*})$.

Therefore, in order to prove Theorem 1, it suffices to show that any optimal
solution of the model defined by (17) attains equality in (10) and (11).

We will prove that equality will be attained in (10) for any optimal solution
in Subsection 3.1 and prove the equality is also attained in (11) for any
optimal solutions in Subsection 3.2.

We make the following assumption:

######  Assumption 1.

We assume that $Mg\Delta H_{1}=0$ and $\beta_{0}=0$. It derives:

$F_{1}\Delta d=0.5M\beta_{1}+\Delta d(A+Bv_{1}+c\beta_{1}).$ (24)

Based on Assumption 1, it is known that $F_{1}\Delta d$ is a positive number
and its reduction caused by a positive value $\beta_{1}$ and $v_{1}$ will not
violate the force and power boundary as defined in (6) and (16). It is known
that $v_{0}=\beta_{0}=0$ and it is also a reasonable assumption that the train
will start from a zero speed.

######  Assumption 2.

We assume that a strict inequality exists for $v_{N}^{2}<\beta_{N}$.

Assumption 2 is used to demonstrate that the model will always find a feasible
solution to keep (11) for $i=1,2,\cdots,N-1$ exact.

######  Remark 1.

If a small variation is introduced to variables $\alpha_{i}$ $v_{i}$ and
$\beta_{i}$, it cannot be guaranteed that the power and force boundaries are
not violated since $F_{i}$ is varied as well. The strategy is then to
introduce a series of small variations and to maintain $F_{n}\Delta
d,\leavevmode\nobreak\ n=2,3,...i$ to remain constant except for $F_{1}\Delta
d$. We can demonstrate that all these small variations, made possible by the
strict equality of relaxation constraints (10) and (11), can eventually reduce
the value of $F_{1}\Delta d$ and thus bring down the objective function
without violating any constraints. Using the proof by contradiction, this will
prove the assumed optimal solution on the strict inequality cannot be the
optimal solution since small variations are possible to be obtained to further
bring down the value of the objective function.

We make the notations as follows:

  * •

$\varepsilon_{i,\alpha}$ to represent a small variation for $\alpha_{i}$ and
$\varepsilon_{i,\alpha}>0$;

  * •

$\varepsilon_{i,\beta}$ to represent a small variation of the original value
of $\beta_{i}$ and $\varepsilon_{i,\beta}>0$;

  * •

$\varepsilon_{i,v}$ to represent a small variation of the original value of
$v_{i}$ and $\varepsilon_{i,v}>0$.

These small variables are all positive and we aim to demonstrate that there is
always a set of $\varepsilon_{i,\alpha}$, $\varepsilon_{i,\beta}$ and
$\varepsilon_{i,v}$ that can bring down the objective function value as long
as the solution is lying on the strict inequality of (10) and (11). We assume
that all the changed variables ($v_{i},\alpha_{i},\beta_{i}$ for
$i=1,2,\cdots,N$) remain positive after the variations are applied.

###  3.1 Proof of exactness: Part 1

Assume that the optimal solution $\alpha_{i}$ and $v_{i}$ for $i=1,2,\cdots,N$
exists with $\alpha_{i}v_{i}>1$. There are small variations
$\varepsilon_{i,\alpha}$ and $\varepsilon_{i,v}$ that make the following
inequality valid.

$(v_{i}-\varepsilon_{i,v})(\alpha_{i}-\varepsilon_{i,\alpha})\geq 1$ (25)

######  Remark 2.

In this case, $\beta_{i}$ remains unchanged while $v_{i}$ is reduced with
$\varepsilon_{i,v}$. This does not violate the constraints
$v_{i}^{2}\leq\beta_{i}$

To make sure $F_{i}\Delta d$ remains unchanged, we reduce the value of
$\beta_{i-1}$ so that $F_{i}\Delta d$ remains constant. By observing (2.2), to
maintain the value of $F_{i}\Delta d$ unchanged, (26) should be applied:

$\varepsilon_{i-1,\beta}=\frac{\Delta dB}{0.5M}\varepsilon_{i,v}$ (26)

where $\varepsilon_{i-1,\beta}$ is also a small variation related to
$\varepsilon_{i,v}$.

To ensure that (11) is valid for the $(i-1)^{th}$ segment, we introduce an
equal reduction on both side of the inequality. (3.1) is applied:

$\displaystyle\beta_{i-1}-\varepsilon_{i-1,\beta}$ $\displaystyle\geq
v_{i-1}^{2}-\varepsilon_{i-1,\beta}$
$\displaystyle=(v_{i-1}-\varepsilon_{i-1,v})^{2}$
$\displaystyle\varepsilon_{i-1,v}$
$\displaystyle=v_{i-1}-\sqrt{v_{i-1}^{2}-\varepsilon_{i-1,\beta}}$
$\displaystyle\varepsilon_{i-1,\beta}$
$\displaystyle=2\varepsilon_{i-1,v}v_{i-1}-\varepsilon_{i-1,v}^{2}.$ (27)

An introduction of $\varepsilon_{i,v}$ can be used to calculate
$\varepsilon_{i-1,\beta}$ based on (26), and $\varepsilon_{i-1,\beta}$ can
then be used to calculate a small variation on $v_{i-1}$, i.e.
$\varepsilon_{i-1,v}$, based on (3.1).

$\varepsilon_{i-1,v}$ and $\varepsilon_{i-1,\beta}$ are both small reduction
introduced to $v_{i-1}$ and $\beta_{i-1}$ so that the model constraints are
not violated. To ensure $F_{i-1}\Delta d$ to remain constant, a small
reduction on $\beta_{i-2}$ needs to be introduced and it is thus defined as
follows.

$0.5M\varepsilon_{i-2,\beta}=(0.5M+C\Delta d)\varepsilon_{i-1,\beta}+B\Delta
d\varepsilon_{i-1,v}$ (28)

This can help calculate the value of $\varepsilon_{i-2,v}$ based on a similar
calculation presented in (3.1) and this process can be iteratively continued
via (3.1) and (3.1) until the first segment is reached.

$\displaystyle\varepsilon_{n,v}$
$\displaystyle=v_{n}-\sqrt{v_{n}^{2}-\varepsilon_{n,\beta}}$ $\displaystyle
n=i-1,i-2,\cdots,1$ (29)

$\displaystyle 0.5M\varepsilon_{n-1,\beta}$ $\displaystyle=(0.5M+C\Delta
d)\varepsilon_{n,\beta}+B\Delta d\varepsilon_{n,v}$ $\displaystyle
n=i-1,i-2,\cdots,2$ (30)

######  Remark 3.

A sufficiently small $\varepsilon_{i,v}$ leads to a set of sufficiently small
$\varepsilon_{n,v},\leavevmode\nobreak\ n=i,i-1,...,1$ and
$\varepsilon_{n,\beta},\leavevmode\nobreak\ \leavevmode\nobreak\ n=i-1,...,1$.

Next we will demonstrate that the change of $v_{n},\leavevmode\nobreak\
n=i-1,...,1$ will not violate the constraint defined by the inequality $1\leq
v_{n}\alpha_{n}$.

######  Remark 4.

Every decrease of $v_{n},\leavevmode\nobreak\ n=i-1,i-2,...,1$ leads an
increase of $\alpha_{n},\leavevmode\nobreak\ n=i-1,i-2,...,1$ to ensure
$v_{n}\alpha_{n}\geq 1$ to remain valid. Due to the constraints defined by
$\sum_{i=1}^{N}\alpha_{i}=T/\Delta d$, we make a special case for $v_{i}$ and
$\alpha_{i}$ so that both of them are reduced. The strict inequality
$1<v_{i}\alpha_{i}$ makes this possible without violating the model
constraints. The reduction of $\alpha_{i}$ will then be used to offset the
increase of $\alpha_{n},\leavevmode\nobreak\ n=i-1,i-2,...,1$.

To calculate $\alpha_{n},\leavevmode\nobreak\ n=i-1,i-2,\cdots,1$, we assume
the relationship applies to the initial optimal solution. This is certainly
valid as the initial optimal solution should be a feasible solution for Model
(17).

$v_{n}\alpha_{n}=D_{n}\geq 1$ (31)

After the introduction of a small variation, the product remains constant
defined as:

$(v_{n}-\varepsilon_{n,v})(\alpha_{n}+\varepsilon_{n,\alpha})=D_{n}$

where, $n=i-1,i-2,\cdots,1$.

Thus, we can obtain the value of $\varepsilon_{n,\alpha},\leavevmode\nobreak\
n=i-1,i-2,\cdots,1$ by

$\varepsilon_{n,\alpha}=\frac{\alpha_{n}\varepsilon_{n,v}}{v_{n}-\varepsilon_{n,v}}.$
(32)

######  Remark 5.

Based on Remarks 3 and 4, we can use the first derivative to evaluate the
relationship between a small variable $v_{n}$ and final total increase of
$\alpha_{n},\leavevmode\nobreak\ n=1,2,\cdots,i-1$.

According to (3.1), the derivative of $\varepsilon_{i-1,v}$ over
$\varepsilon_{i-1,\beta}$ is given as follows:

$\frac{d(\varepsilon_{i-1,v})}{d(\varepsilon_{i-1,\beta})}=\frac{1}{\sqrt{v_{i-1}^{2}-\varepsilon_{i-1,\beta}}}>0$

By observing (3.1), it is known that this is a strictly increasing monotonic
function over either $\varepsilon_{i-1,\beta}$ and $\varepsilon_{i-1,v}$

By observing (32), we are able to obtain its first derivative of
$\varepsilon_{n,\alpha}$ over $\varepsilon_{n,v}$ as follows:

$\frac{d(\varepsilon_{n,\alpha})}{d(\varepsilon_{n,v})}=\frac{\alpha_{n}}{v_{n}-\varepsilon_{n,v}}+\frac{\alpha_{n}\varepsilon_{n,v}}{(v_{n}-\varepsilon_{n,v})^{2}}=\frac{D_{n}}{(v_{n}-\varepsilon_{n,v})}>0$

The reduction of $\alpha_{i}$ will be equal to the sum of all
$\varepsilon_{n,\alpha},\leavevmode\nobreak\ n=i-1,i-2,\cdots,1$, represented
by:

$\varepsilon_{i,\alpha}=\sum_{n=1}^{i-1}\varepsilon_{n,\alpha}.$ (33)

Given that each $\varepsilon_{n,\alpha}$ is a strictly increasing monotonic
function over $\varepsilon_{n,v}$ which is also a strictly increasing
monotonic function over $\varepsilon_{i,v}$, by the property of monotonic
functions, it can be deducted that $\varepsilon_{n,\alpha}$ is a strictly
increasing monotonic function over $\varepsilon_{i,v}$.

Therefore, there exist sufficiently small numbers $\varepsilon_{n,v}$ and
$\varepsilon_{n,\alpha}$ that reduce the value of the assumed optimal solution
$v_{i}$ and $\alpha_{i}$ which are constrained by a strict inequality
$1<\alpha_{i}v_{i}$ to generate a new feasible solution $\tilde{v}_{n}$ and
$\tilde{\alpha}_{n}$, defined by (34) and (35) respectively:

$\tilde{v}_{n}=v_{n}-\varepsilon_{n,v},\leavevmode\nobreak\ n=i,i-2,\cdots,1$
(34)

$\tilde{\alpha}_{n}=\begin{cases}\alpha_{n}-\varepsilon_{n,\alpha}\text{ if
}\leavevmode\nobreak\ n=i\\\ \alpha_{n}+\varepsilon_{n,\alpha}\text{ if
}\leavevmode\nobreak\ n=i-1,i-2,\cdots,1\end{cases}$ (35)

Figure 1: An illustrative graph for the part 1 of exactness proof: a small
variation of $v_{i}$ leads to a small variation of $\alpha_{i}$ without
violating constraints and finally reduce the objective function value.
Relevant applied equation number for calculations is listed in the figure.

Fig. 1 shows an illustration for the exactness proof in part 1. On the
segments other than the first one, $F_{i}\Delta d$ remains constant. On the
first segment, the initial value of $\beta_{1}$ and $v_{1}$ will be reduced
with a small value $\varepsilon_{1,\beta}$ and $\varepsilon_{1,v}$ and this
does not violate the power and force boundaries and other constraints as
claimed in Assumption 1.

In other words, the new set of variables is the feasible solution and attains
a lower value of the objective function. This contradicts the initial
assumption that $v_{i}$ and $\alpha_{i}$ are the optimal solution. This shows
that there is no such a solution $v_{i}$ and $\alpha_{i}$ that is on the
strict inequality $1<\alpha_{i}v_{i}$ and also achieve the optimal solution
for Model A. Model B will always achieve its optimal solution on equality
constraint $1=\alpha_{i}v_{i}$. This completes Part 1 of the proof of
exactness.

###  3.2 Proof of exactness: Part 2

Similarly, we seek similar sets of small variations which always prevent the
variables on the strict inequality from being optimal. First, assume that the
optimal solution exists on the strict inequality of (11), and thus there are
two small variations that make following constraint valid:

$(v_{i}+\varepsilon_{i,v})^{2}\leq\beta_{i}-\varepsilon_{i,\beta}$

In Fig. 2, we illustrate the part 2 of exactness proof for constraint 11.

Figure 2: An illustrative graph for the part 2 of exactness proof: with an
assumed relaxed constraints at the $N^{th}$ segment, we are able to achieve a
feasible solution using a set of small variations and finally reduce the value
of the objective function. The yellow arrow within each block shows if the
variation is to increase or to decrease the original variable and the relevant
applied equation for calculation of different variables is included from one
variable block to another.

We introduce a small variation $\varepsilon_{i-1,\beta}$ to bring down the
initial optimal speed $v_{i-1}$ and it can be demonstrated that using (3.1)
and (3.1), we can bring down the value of initial value of $\beta_{n}$ and
$v_{n}$ where $n=i-1,i-2,\cdots,1$ without violating the constraints.

In addition, by adopting (32), it is possible to calculate the increase of
$\alpha_{n},\leavevmode\nobreak\ n=i-1,i-2,\cdots,1$ and the value of the
reduction of $\alpha_{i}$ can be calculated by (33). Based on (33), it is
possible to calculate the small increase of $v_{i}$, i.e. $\varepsilon_{i,v}$.
To maintain a constant value of $F_{i}\Delta d$, based on (2.2), the following
two relations need to be applied.

$\varepsilon_{i,\beta}(C\Delta d+0.5M)=B\Delta
d\varepsilon_{i,v}+0.5M\varepsilon_{i-1,\beta}$ (36)

$\varepsilon_{N,\beta}(C\Delta d+0.5M)=B\Delta
d\varepsilon_{N,v}+0.5M\varepsilon_{N-1,\beta}$ (37)

Based on Assumption 2, there exists a small gap between $\beta_{N}$ and
$v_{N}$ and a small variation can be applied to reduce the value of
$\beta_{N}$ and $\beta_{N}-1$ and increase the value of $v_{N}$.

According to (3.1) and (3.1), a small variation $\varepsilon_{N-1,\beta}$ on
$\beta_{N-1}$ gives rise to a series of reduction on
$v_{n},\leavevmode\nobreak\ n=N-1,N-2,\cdots,i+1$ and based on Remark 5, it is
known that $\sum_{n=i+1}^{N-1}\alpha_{n}$ is strictly increasing monotonic
over $\varepsilon_{N-1,\beta}$. Since the speed $v_{n}$ is reducing,
$\alpha_{n}$ is increasing. On the other hand, $\varepsilon_{N,\alpha}$ is a
strictly increasing monotonic function of $\varepsilon_{N,v}$. Note that
$v_{N}$ is increasing with a variation of $\varepsilon_{N,v}$.

Now define a function $f(\varepsilon_{N-1,\beta})$ as follows:

$f(\varepsilon_{N-1,\beta})=\sum_{n=i+1}^{N-1}\alpha_{n}-\varepsilon_{N,\alpha}$
(38)

This function is a continuous function since no extreme value is taken and has
a feasible domain of $[0,\bar{\varepsilon}_{N-1,\beta}]$ where
$\bar{\varepsilon}_{N-1,\beta}$ is determined by (37) when $\varepsilon_{v,N}$
takes a value of zero, as shown below:

$\bar{\varepsilon}_{N-1,\beta}=\left(1+\frac{\Delta
dC}{0.5M}\right)\varepsilon_{N,\beta}.$

On the boundary of the feasible domain, (38) takes values of
$-\varepsilon_{N,\alpha}<0$ and $\sum_{n=i+1}^{N-1}\alpha_{n}>0$. Based on
Intermediate Value Theorem, there exists one value of
$\varepsilon_{N-1,\beta}$ that causes the function to take a value of zero. We
use $\tilde{\varepsilon}_{N-1,\beta}$ and $\tilde{\varepsilon}_{i,\beta}$ to
denote the corresponding values that leads (38) to take a value of zero. With
the existence of $\tilde{\varepsilon}_{N-1,\beta}$ and
$\tilde{\varepsilon}_{i,\beta}$, it is known that every $F_{n}\Delta d$ for
$n=i+1,i+2,\cdots,N$ will be maintained constant and the constraints on
$\alpha_{n}$ for $n=i+1,i+2,\cdots,N$ remain valid as the increase of
$\alpha_{n}$ for $n=i+1,i+2,\cdots,N-1$ is offset by the decrease of
$\alpha_{N}$.

Given a sufficiently small variation of $\tilde{\varepsilon}_{i,\beta}$, we
are able to conduct similar analysis based on Intermediate Value Theorem and
(36) to show that there exists $\tilde{\varepsilon}_{i-1,\beta}$ and
$\tilde{\varepsilon}_{i,v}$ that makes the constraints on
$\alpha_{n}\leavevmode\nobreak\ \leavevmode\nobreak\ n=1,2,\cdots,i$ valid and
meanwhile brings down a series of reduction on $v_{n}$ and $\beta_{n}$ for
$n=1,2,\cdots,i-1$. Based on Assumption 1, the existence of variations will
decrease the objective function on the first segment and this proves that (11)
remain exact for any index number $n=1,2,\cdots,N-1$. It is worth noting that
for the case of the $N^{th}$ segment, the $\beta_{N}$ will always keep
approaching $v_{N}^{2}$ and virtually this constraint remains exact even with
the assumption on its strict inequality in Assumption 2.

In a summary, we have demonstrated that all optimal solution of Model B will
only be obtained on the equality of relaxation constraints defined by (10) and
(11). This thus proves that Theorem 1 is true based on Assumption 1 and 2. As
a matter of fact, as long as there is any $F_{n}\Delta d$ where $n=1,2,\cdots
i$ where a small reduction of $F_{n}\Delta d$ does not violate the power and
force boundary defined by (16) and (6), e.g. $F_{n}\Delta d>0$, the proof
remains valid.

##  4 Numerical verification on a case with a steep gradient

In this section, we demonstrate a quick numerical verification on the model
exactness using a case study with a steep gradient.

The modeling parameters are listed in Table 1.

Table 1: Modeling parameters of the traction system Parameter |  Value
---|---
$T(\rm s)$ |  260
$M(\rm t)$ |  144
$F_{t,max}(\rm kN)$ |  230.81
$F_{b,max}(\rm kN)$ |  230.81
$P_{t,max}(\rm kW)$ |  2520
$P_{b,max}(\rm kW)$ |  2520
$\eta_{t}$ |  0.9
$\eta_{b}$ |  0.6
$A(\rm kN)$ |  3.0016
$B(\rm kN/(km/h))$ |  2.016e-2
$C(\rm kN/(km^{2}/h^{2}))$ |  6.9692e-4
Figure 3: Optimal speed trajectory and tractive/braking efforts for a typical
urban rail case based on convex optimization. (a) It demonstrates the optimal
speed trajectory with varying speed limit and steep gradient. (b) The
tractive/braking effort shows that the train optimal operation consists of
maximum traction/braking, partial traction/braking, and coasting.

In Fig. 3, we present a typical urban rail case with an arbitrary steep
downhill gradient between 3000 m and 4000 m on which the train needs to brake
to maintain a constant speed. On distance 4000 m, the change of the gradient
has been counteracted by the varying efforts. The results show that the model
can be solved and quickly achieve an optimal solution with a steep gradient.

Figure 4: Comparison of variables. (a) Comparison between the reciprocal of
the candidate speed $1/v_{i}$ and $\alpha_{i}$. (b) Comparison between the
square of the candidate speed $v_{i}^{2}$ and the variable $\beta_{i}$. Both
figures show that when the optimization model achieves the optimal solution,
inequality constraints (10) and (11) achieves equality.

Fig. 4 shows a comparison of two sets of variables of interest and
demonstrates that equality for (10) and (11) will be attained when the optimal
solution is achieved. More numerical results and discussion can be found in
our newly published preprint paper[11].

##  5 Conclusions

As a continued work of [11], this paper presents the exactness proof for an
EETC model based on convex optimization referred to as Model B. We applied two
constraint relaxation to convert an non-convex Model, i.e. Model A, into a
convex Model B. The convexity of Model B was first verified and the exactness
proof of Model B was then conducted. With its convexity and relaxation
exactness, the proposed model with mature solution approaches such as
interior-point barrier methods boast a very high computational efficiency and
global optimality compared to its counterpart such as MILP, Psuedospectral
Method and other non-linear programming methods. Its high flexibility could
enhance the model’s adaptivity for other types of transportation modes to
further reduce the energy consumption.

## References

  * [1] E. Khmelnitsky.  On an optimal control problem of train operation.  IEEE Transactions on Automatic Control, 45:1257–1266, 7 2000.
  * [2] P Howlett.  The optimal control of a train.  Annals of Operations Research, 98:65–87, 2000.
  * [3] Rongfang Liu and Iakov M Golovitcher.  Energy-efficient operation of rail vehicles.  Transportation Research Part A: Policy and Practice, 37(10):917–932, 2003.
  * [4] Amie Albrecht, Phil Howlett, Peter Pudney, Xuan Vu, and Peng Zhou.  The key principles of optimal train control—part 1: Formulation of the model, strategies of optimal type, evolutionary lines, location of optimal switching points.  Transportation Research Part B: Methodological, 94:482–508, 2016\.
  * [5] Amie Albrecht, Phil Howlett, Peter Pudney, Xuan Vu, and Peng Zhou.  The key principles of optimal train control—Part 2: Existence of an optimal strategy, the local energy minimization principle, uniqueness, computational techniques.  Transportation Research Part B: Methodological, 94:509–538, 2016\.
  * [6] Rob M.P. Goverde, Gerben M Scheepmaker, and Pengling Wang.  Pseudospectral optimal train control.  European Journal of Operational Research, 292:353–375, 7 2021.
  * [7] Yihui Wang, Bart De Schutter, Ton J.J. van den Boom, and Bin Ning.  Optimal trajectory planning for trains - a pseudospectral method and a mixed integer linear programming approach.  Transportation Research Part C: Emerging Technologies, 29:97–114, 2013.
  * [8] Pengling Wang and Rob M P Goverde.  Multiple-phase train trajectory optimization with signalling and operational constraints.  Transportation Research Part C: Emerging Technologies, 69(Supplement C):255–275, 2016.
  * [9] Hongbo Ye and Ronghui Liu.  Nonlinear programming methods based on closed-form expressions for optimal train control.  Transportation Research Part C: Emerging Technologies, 82:102–123, 2017.
  * [10] Chaoxian Wu, Wenrui Zhang, Shaofeng Lu, Zhaoxiang Tan, Fei Xue, and Jie Yang.  Train Speed Trajectory Optimization With On-Board Energy Storage Device.  IEEE Transactions on Intelligent Transportation Systems, 20(11):4092–4102, nov 2019.
  * [11] Minling Feng, Kunpeng Wu, and Shaofeng Lu.  A fast-solved model for energy-efficient train control based on convex optimization, 2022.
  * [12] Stephen P. Boyd and Lieven. Vandenberghe.  Convex optimization.  Cambridge University Press, 2004.