text:
\section{INTRODUCTION}
\label{sec:intro}
Industrial robots are increasingly deployed in applications such as spray coating \cite{spray_coating}, arc welding \cite{arc_welding}, deep rolling \cite{deep_rolling}, surface grinding \cite{surface_grinding}, cold spraying \cite{cold_spray}, etc., where the tool center point (TCP) frame attached to the end effector needs to track complex geometric paths in both position and orientation. In most applications, the task performance is characterized by the motion speed (how long it takes to complete the task), motion uniformity (how much the speed varies along the path), and motion accuracy (the maximum position and orientation tracking errors along the path).
There are generally two ways to program industrial robots for the motion tracking task: (1) The motion primitive method that uses  vendor-specific proprietary robot programming languages (e.g., \cite{Inform_manual,RAPID_manual,KAREL_manual,Vplus_manual,VAL3_manual,KRL_manual}) consisting of a sequence of pre-defined motion primitives, and (2) The streaming setpoint method which uses an external command mode to stream the desired joint position as setpoints to the low level joint servo controller (e.g., \cite{MotoPlus,EGM,LLI,Kuka_open}).
The motion primitive method is far more widely used in industries.  It  decomposes the desired motion with predefined motion primitives consisting of waypoints connected by motion segments.  The streaming setpoint method is not limited to a small set of motion primitives, but it typically results in poorer tracking accuracy due to the lower sampling rate
and the additional filtering and latency.

To program an industrial robot to follow a complex path using motion primitives, the most common approach is to use it like a machine tool \cite{Robodk_sim}, focusing only on the TCP motion. There are several advanced offline programming software (e.g., \cite{Robodk,RobotMaster,Octobpuz}) that convert a dense set of TCP waypoint
to the robot program of a given robot vendor.
The actual robot motion from the execution of the robot program will depend on the motion primitives and their parameters (waypoint locations, commanded path speed along the motion segment, size of the blending zone between motion segments) which are affected by the robot joint servo controller and the robot joint motion constraints (joint velocity and acceleration limits).  Robot controller and robot joint constraints are largely unknown to the user.  Therefore, programming a robot to achieve the desired performance is currently a time-consuming and largely manual exercise in tuning the motion primitive parameters and the result may be far from optimal.  Compounding the challenges are the impact of additional parameters: a redundant
roll degree of freedom in the tool frame, relative pose between
the target curve and the robot base, and the robot configuration
(out of the 8 possible choices in a 6 degrees of freedom (DoF) revolute robot).

The streaming setpoint method avoids motion blending, but the actual robot motion is difficult to predict, again because of the unknown low level controller and robot motion constraints, in addition to the latency and filtering present in the streaming operation.
Motion planning tools such as Tesseract \cite{Tesseract} (based on TrajOpt \cite{Trajopt}), tries to optimize the robot joint path for path speed subject to tracking error and speed uniformity constraints.  Such optimization involves a large number of variables and is plagued by computation time, presence of local minima, and the lack of accurate prediction of the actual robot motion.

\begin{figure}[h!]
    \centering
    \includegraphics[width=0.30\textwidth]{images/fig1x.pdf}
    \caption{Illustration of the motion primitive method for path following.}
    \label{fig:f1}
\end{figure}

A key motivation of this work is the \emph{cold spray} application which impinges particles onto a surface at high speed to form protective coating \cite{cold_spray}.  For such applications, the speed of the TCP needs to be \emph{high and uniform} to guarantee fast execution and uniform distribution of the deposited material.
As shown in Fig.~\ref{fig:f1}, the spraying path is given as a dense sequence of points $\{\mathbf{p}^\ast_i\}_{i=1}^K$ for the TCP position and the spraying direction in the negative tool $z$-axis is given by the surface normal $\{\mathbf{n}^\ast_i\}_{i=1}^K$.   Note that these data specify 5-DoF of the TCP at each point, with the remaining free DoF being the rotation about the tool $z$-axis.  Once the TCP frame is fixed, there are multiple corresponding robot joint configurations (inverse kinematics solutions).
We also have an additional 6 free DoF in the position and orientation of the curve, parameterized by the position of the center of mass $\mathbf{p}_{curve}$ and the angle-product representation of the orientation $\bm{\beta}_{curve}$.
The motion primitive method first chooses all these free parameters, we call this step {\em redundancy resolution}.  The TCP frames are then converted to motion primitives as shown in Fig.~\ref{fig:f1} for execution by the robot controller.
This paper presents a new approach to optimizing the motion primitives by decomposing the problem into a series of subproblems. We first use a combination of evolution algorithm and local optimization to resolve redundancy. Then we use a greedy algorithm for fitting the curve with the least number of motion primitives. Finally, an iterative gradient descent algorithm updates the waypoints to minimize the worst case tracking error.  To evaluate this approach, the industry co-authors of this work suggested two test curves, curve one contains a range of high curvature moves and curve two is the leading edge of a sample turbine fan blade.  A baseline approach is developed based on the current industry practice.  The evaluation metric is the average path speed subject to the specified trajectory tracking and path speed uniformity constraints.  We implemented the motion primitive optimization approach on an ABB IRB 6640 robot in simulation (using RobotStudio) and physical testbed, and on an FANUC M710iC-70 robot in simulation (using RoboGuide).  In all cases, the performance vastly improves over the baseline, ranging over 200\% to 300\%.

\noindent \textbf{Notation:} Throughout the paper, $\mathbf{p}\in \mathbb{R}^3$ denotes the Cartesian position of the TCP, $(v,\omega)$ denotes the linear and angular velocity of TCP, $p_0
\in \mathbb{R}^3$ denote the TCP position at the robot zero configuration, $\bm{\beta} \in \mathbb{R}^3$ denotes the angle-product representation of the tool frame orientation,
$\mathbf{q}\in \mathbb{R}^6$ denotes the joint position of the robot. The function $\mathtt{fwd}:\mathbf{q}\rightarrow(\mathbf{p},\bm{\beta})$ denotes the forward kinematic map. We also use the notations $\mathbf{p}(\mathbf{q})$ and $\bm{\beta}(\mathbf{q})$ to represent the forward kinematics. The notation $e_z(\mathbf{q}) \in \mathbb{R}^3$ is a unit vector denoting the $z$-axis of the TCP in the robot base frame.
The TCP spatial velocity is denoted by $\begin{bmatrix} \omega \\ v \end{bmatrix} = J(\mathbf{q}) \dot{\mathbf{q}} $, where $J$ is the Jacobian matrix.

\begin{tcolorbox}\textbf{Contribution:} We present a method for calculating a set of motion primitives and their respective parameters (see Sec. \ref{sec:primitives} for more details) so that the TCP tracks the desired positions and orientations with sufficient accuracy and speed uniformity (see Sec. \ref{sec:spec} for more details) as fast as possible.
\end{tcolorbox}

\section{Motion Primitives and Their Executions}
\label{sec:primitives}

\subsection{Motion Primitives}
Generally all robot programming languages support three types of motion primitives: MoveL, MoveC and MoveJ, which are described below.
\begin{itemize}
    \item MoveL:  The TCP moves linearly in Cartesian space and linearly in angle-product orientation from current pose to desired pose. This command is parameterized by the initial and final TCP positions and angle-product orientations.
    \item MoveC:  The TCP moves on a circular arc in Cartesian space and linearly in angle-product orientation from current pose to desired pose. This command is parameterized by the initial and final TCP positions, an intermediate TCP position (to define the circular arc), and the initial and final angle-product orientations, an intermediate orientation (to define rotation orientation).
    \item MoveJ:  The robot moves linearly in the \emph{joint angles space}. Thus, the TCP path may not be linearly nor circular. This command is parametrized by the initial and final joint angles.
\end{itemize}

\subsection{Blending}
Given a series of motion primitive segments, the desired tool path is  blended together to avoid sharp turns \cite{6-5_spline}. The size of the blending zone is specified with the primitives. Larger blending zone means smoother transition to next segment, and hence higher and more uniform speed, but at the cost of large tracking error at the waypoint (intersection of motion segments). Small blending zone improves the tracking accuracy but could lead to sharp corner requiring slower speed with larger variation to meet the motion constraints. In current practice, expert human operators perform extensive manual adjustments of the waypoints, blending regions, and motion segment commanded speeds to ensure that tracking accuracy and cycle time requirements are both met.  In our approach, we find the smallest uniform blending zone that avoids excessive speed variation around waypoints to balance between the two objectives.

\section{Specifications and Proposed Approach}
\label{sec:spec}

Based on the cold spray application, the trajectory tracking performance specification is given by: \\
\noindent \textbf{1. Positional tracking accuracy:} the maximum positional tracking error of the TCP is less than 0.5 mm.\\
\textbf{2. Orientation tracking accuracy:} the maximum orientation tracking error of tool $z$-axis is less than 3$^\circ$.\\
\textbf{3. Speed uniformity:} the standard deviation of the TCP speed is less than 5\%.\\

Our approach decomposes the motion primitive optimization problem into three steps as shown in Fig. \ref{fig:flowchart}. \\
{\bf 1. Redundancy Resolution}: Optimize the redundant DoF (curve pose: ($\mathbf{p}_{curve}$,$\bm{\beta}_{curve}$)), the redundant DoF along the path, and the robot arm configuration. The output of this process is the full 6 DoF reference path for the TCP and the corresponding joint path. \\
{\bf 2. Motion Primitive Planning}: Generate a sequence of motion primitives (waypoints, blending zones, and path speed) to track the reference path. \\
{\bf 3. Waypoint Adjustment}: The motion sequence is executed in simulations or experiments.  Tracking error from the executions is used to optimize the motion primitives parameters. \\
These steps are described in greater details below.

\begin{figure}[h]
    \centering
    \includegraphics[width=0.48\textwidth]{images/flowchart.jpg}
    \caption{Proposed Workflow for Spatial Curve Tracking with Motion Primitives.
    }
    \label{fig:flowchart}
\end{figure}

\subsection{Redundancy Resolution}
In contrast to computer numerical control (CNC) machines, robot pose significantly affects the TCP motion due to the robot joint constraints, including the joint position, velocity, and acceleration limits.  There are multiple redundant DoF that may be exploited to place the robot in a more advantageous configuration without compromising the TCP motion:
\begin{itemize}
    \item Tool orientation: Rotation around the tool $z$-axis at each trajectory point.
    \item Curve pose: Position $\mathbf{p}_{curve}$ and orientation $\bm{\beta}_{curve}$ (6-doF).
    \item Robot pose: One out of maximum eight joint configurations due to the non-uniqueness of the inverse kinematic map.
\end{itemize}
Our approach is to first resolve the tool redundant DoF by finding the complete robot joint path based on a given curve pose and robot pose.  At each point on this joint path, we can find the path velocity limit based on the joint velocity and acceleration limits.  We then choose the curve pose and robot pose to maximize the lowest path velocity limit.  The joint acceleration limit is due to the motor torque constraint.  Since the arm inertia is configuration dependent, the joint acceleration limit is also configuration dependent (hence, an outstretched arm has a lower acceleration limit).  We will show how to estimate this limit through simulation.

\subsubsection{Tool Orientation Resolution}
\label{sec:toolorientation}

For a given part position and orientation, $(\mathbf{p}_{curve},\bm{\beta}_{curve})$, and robot pose, we have
$\{\mathbf{p}^\ast_i, \mathbf{n}^\ast_i\}_{i=1}^K$ in the robot base frame and the initial joint angle $\mathbf{q}_0$.
We select the redundant DoF so that the TCP motion remains on the path with minimal incremental joint motion.  Since joint motion and TCP motion are related through the arm Jacobain $J(\mathbf{q})$, this procedure tends to steer the robot away from singularities without compromising TCP tracking.  We pose the incremental optimization as a quadratic program (QP):
Given $\mathbf{q}_{k-1}$, $k=1,\ldots,K$, find the incremental motion $\delta q$ to move to the next point $(\mathbf{p}^\ast_k, \mathbf{n}^\ast_k)$ without excessive joint motion and within the joint limits:
\begin{equation}
\begin{aligned}
    &\min_{\delta q} \norm{J_r(\mathbf{q}_{k-1}) \delta q -\nu^*}^2 + W_q\norm{\delta q}^2 \\
    &\mbox{subject to } q_{\min} \preccurlyeq \mathbf{q}_{k-1}+\delta q \preccurlyeq q_{\max}
\end{aligned}
\label{eq:QP_reduced}
\end{equation}
where $(q_{min},q_{max})$ are joint limits, $W_q$ is a weightfactor,  $\nu^*$ is the TCP position and surface normal increment, and $J_r$ is the reduced Jacobian mapping $\delta q$ to the position and surface normal increment:
\begin{equation}
    \nu^* := \begin{bmatrix} \mathbf{e_z}(\mathbf{q}_{k-1})-\mathbf{n}^*_{k} \\ \mathbf{p}(\mathbf{q}_{k-1})-\mathbf{p}^*_{k} \end{bmatrix},
\end{equation}
\begin{equation}
    J_r(\mathbf{q}):= \begin{bmatrix} -\mathbf{e_z}(\mathbf{q})^\times & 0 \\ 0 & I \end{bmatrix} J(\mathbf{q}).
\end{equation}
The next joint position is simply the current joint position incremented by $\delta q$.  Since the Jacobian mapping is only approximate for finite increments, we optimize the increment size by writing $\mathbf{q}_k = \mathbf{q}_{k-1}+\alpha \delta q$ and choose $\alpha$ from a line search:
\[
\min_\alpha  \norm{\begin{bmatrix} \mathbf{e_z}(\mathbf{q}_{k})-\mathbf{n}^*_{k} \\ \mathbf{p}(\mathbf{q}_{k})-\mathbf{p}^*_{k} \end{bmatrix}}.
\]
The resulting sequence of joint angles  $\{\mathbf{q}_k\}_{k=0}^K$ now completely specifies the TCP position and orientation.

\subsubsection{Trajectory Traversal Speed Estimation}

The robot TCP speed limit is determined by its joint velocity and acceleration limits.  Industrial robots provide joint velocity limits in the data sheets \cite{6640_manual}, but the acceleration limit is due to the torque limit and is therefore affected by the robot dynamics, load, and robot configuration. Since the dynamical models of industrial robots are typically unknown, we
 establish a rough estimation by using vendor-provided dynamical simulation such as the ABB RobotStudio.

 Since the arm inertia is most heavily influenced by the shoulder and elbow joints (joints 2 and 3), we assume that the spherical wrist joints (joint 4, 5, and 6) have constant joint acceleration limits, while joints 1, 2 and 3 acceleration limits depend on $(q_2,q_3)$. The recorded result is shown in Table~\ref{table:6640_limit} and Fig.~\ref{fig:j_acc}.

\begin{table}[h!]
\centering
\begin{tabular}{||c c c c c c c||}
 \hline
 Joint & 1 & 2 & 3 & 4 & 5 & 6 \\ [0.2ex]
 \hline\hline
 $\dot q_{max}$ ($rad/s$)&   1.745 & 1.571 & 1.571 & 3.316 & 2.443 & 3.316   \\
 $\ddot q_{max}$ ($rad/s^2$)&       *   & *     & *     & 42.5 &  36.8 & 50.5   \\[0.5ex]
 \hline
\end{tabular}
\caption{ABB IRB6640 joint velocity and acceleration limits. * = configuration dependent, see Fig. \ref{fig:j_acc}.}
\label{table:6640_limit}
\end{table}

\begin{figure}[h]
     \centering
     \begin{subfigure}[b]{0.22\textwidth}
         \centering
         \includegraphics[width=\textwidth]{images/j1_acc.png}
         \caption{Joint1}
         \label{fig:j1_acc}
     \end{subfigure}
     \begin{subfigure}[b]{0.22\textwidth}
         \centering
         \includegraphics[width=\textwidth]{images/j2_acc.png}
         \caption{Joint2}
         \label{fig:j2_acc}
     \end{subfigure}
     \begin{subfigure}[b]{0.22\textwidth}
         \centering
         \includegraphics[width=\textwidth]{images/j3_acc.png}
         \caption{Joint3}
         \label{fig:j3_acc}
     \end{subfigure}
        \caption{Configuration Dependent Joint Acceleration Limit: Recorded maximum joint acceleration of joint 1,2,3 with respect to different joint 2 and 3 configuration.}
        \label{fig:j_acc}
\end{figure}
\def\mathbf{q}{\mathbf{q}}
\def\mathbf{\dot q}{\mathbf{\dot q}}
\def\mathbf{\ddot q}{\mathbf{\ddot q}}
\def\mathbf{p}{\mathbf{p}}
The joint acceleration limits are stored as a look-up table dependent on $(q_2,q_3)$: $\mathbf{\ddot q}_{\max}(q_2,q_3)$. For each robot joint angle point $\mathbf{q}_k$ along the path (from Sec.~\ref{sec:toolorientation}), the maximum joint velocity limit is given by
\begin{equation}
    \mathbf{\dot q}_{\max_k} = \min\left(\mathbf{\dot q}_{\max},\mathbf{\dot q}_{k-1}+\Delta t_k \mathbf{\ddot q}_{\max}(q_2,q_3) \right).
\end{equation}
For a specified (scalar) path speed $v_d$, $\mathbf{p}_k = \mathbf{p}_{k-1} + v_d \Delta t_k$.  Therefore,
\begin{equation}
    \Delta t_k = \frac{\mathbf{p}_k-\mathbf{p}_{k-1}}{v_d}.
\end{equation}
Hence the path velocity limit for point $k$ on the path is
\begin{equation}
    v_k = \norm{J(\mathbf{q}_k){\mathbf{q}_{\max}}_k}.
    \label{eq:vk}
\end{equation}

\subsubsection{Differential Evolution with Optimizing Parameters}

\begin{figure}[h]
     \centering
     \begin{subfigure}[b]{0.2\textwidth}
         \centering
         \includegraphics[width=\textwidth]{images/curve1_baseline_pose.png}
         \caption{Curve1 Baseline Pose}
         \label{fig:curve1_baseline_pose}
     \end{subfigure}
     \begin{subfigure}[b]{0.2\textwidth}
         \centering
         \includegraphics[width=\textwidth]{images/curve1_opt_pose.png}
         \caption{Curve1 Optimized Pose}
         \label{fig:curve1_opt_pose}
     \end{subfigure}
     \begin{subfigure}[b]{0.2\textwidth}
         \centering
         \includegraphics[width=\textwidth]{images/curve2_baseline_pose.png}
         \caption{Curve2 Baseline Pose}
         \label{fig:curve2_baseline_pose}
     \end{subfigure}
     \begin{subfigure}[b]{0.2\textwidth}
         \centering
         \includegraphics[width=\textwidth]{images/curve2_opt_pose.png}
         \caption{Curve2 Optimized Pose}
         \label{fig:curve2_opt_pose}
     \end{subfigure}
        \caption{Baseline curve pose (left) vs optimized curve pose (right) for Curve 1 and Curve 2 (described in Sec. \ref{sec:curves}.)}
        \label{fig:pose_opt}
\end{figure}

With the path speed limit parameterized by the curve pose and robot pose, we can optimize them to maximize the lowest speed limit:  ($\min(\{\mathbf{v}_k\}_{k=0}^K)$) could be formed as
\begin{equation}
\max_{\mathbf{p}_{curve},\bm{\beta}_{curve},\mathbf{q}_0} \min_k v_k
\label{eq:diff_evo}
\end{equation}
where $v_k$ is given by \eqref{eq:vk}. We use the differential evolution algorithm \cite{diff_evo}
to find the global optimizing solution for two example curves.  The optimized poses are as shown in Fig.~\ref{fig:pose_opt} versus the baseline (see Section~\ref{execution_baseline}).

\subsection{Motion Primitive Planning with Greedy Fitting}

Existing robot CAD/CAM software only use MoveL with a specified spatial curve. In order to take advantage of MoveC and MoveJ, we have developed a spatial curve (represented in joint space and Cartesian space) greedy fitting. The fitting process is performed sequentially along the trajectory, starting from the initial point and bisection search for the longest possible primitive type within the error threshold. For the initial segment, all regressions are unconstrained, while subsequent sequences are subject to continuity constraint (regression must pass through last segment endpoint).
Fig.~\ref{fig:fitting} shows a visualization of 30 equally spaced MoveL's and Greedy Fitting with all three primitives on a test curve. While 30 linear segments will result in 1.47mm max position tracking error, greedy fitting only uses 17 segments to achieve a max position tracking error under 0.3mm.

\begin{figure}[h]
     \centering
     \begin{subfigure}[b]{0.23\textwidth}
         \centering
         \includegraphics[width=\textwidth]{images/30L_curve1.png}
         \caption{Equally spaced MoveL's}
         \label{fig:eq_movel}
     \end{subfigure}
     \begin{subfigure}[b]{0.23\textwidth}
         \centering
         \includegraphics[width=\textwidth]{images/greedy.jpg}
         \caption{Greedy fitting}
         \label{fig:greedy}
     \end{subfigure}
        \caption{Curve fitting visualization.}
        \label{fig:fitting}
\end{figure}

\subsection{Waypoints Adjustment}
\label{waypoint_adjust}

The execution of the robot program based on the motion primitives would likely result in a trajectory that does not meet the performance specification. We can iteratively modify the waypoints to reduce the trajectory tracking error using a descent algorithm.   We use the following two algorithms:
\begin{itemize}
    \item Error Compensation: Push all waypoints in the error direction.
    \item Multi-peak Model Based Gradient Descent: Use an approximate trajectory
    to calculate the numerical gradient and adjust waypoints based on the gradient descent. To simplify computation, we only apply to points at the error peaks.
\end{itemize}
Error compensation takes place in the first few iterations, aiming to bring down the tracking error at each waypoint while keeping the same primitive choices: $\begin{bmatrix} \mathbf{p}_{bp,j} \\ \bm{\beta}_{bp,j} \end{bmatrix}=\begin{bmatrix} \mathbf{p}_{bp,j} \\ \bm{\beta}_{bp,j} \end{bmatrix}+\gamma E(\begin{bmatrix} \mathbf{p}_{b,j} \\ \bm{\beta}_{bp,j} \end{bmatrix})$, where $E(\begin{bmatrix} \mathbf{p}_{bp,j} \\ \bm{\beta}_{bp,j} \end{bmatrix})$ is the error vectors at
$j_{th}$ waypoint and $\gamma$ is the step size.

Error compensation step only focuses the error at the waypoints, while the tracking error between waypoints may be unacceptably large.  We next apply a gradient descent method using an approximate numerical gradient to reduce the largest tracking errors.  For a given set of waypoints, we can efficiently generate a predicted joint trajectory using cubic spline interpolation at each waypoint. This predicted trajectory may then be used to find the numerical gradient by perturbing the waypoints and record the corresponding tracking error.  As the trajectory error is predominately affected by its neighboring waypoints, we only generate a tri-diagonal gradient, and only for the peak trajectory errors.

\section{Implementation and Evaluation}
\subsection{Test curves}
\label{sec:curves}

We use two representative curves. Curve 1, shown in Fig. \ref{fig:fitting}, is a multi-frequency sine curve on a parabolic surface. This curve is representative of a high curvature spatial curve. Curve 2, shown in Fig. \ref{fig:pose_opt}, is extracted from the leading edge of a generic fan blade model, which is a typical case for cold spraying applications in industry.

\subsection{Baseline}
\label{execution_baseline}
In consultation with industry experts, we establish a baseline performance based on the current practice.  The curve is placed in the middle of the robot workspace, away from singularities. The tool redundancy is resolved by choosing the tool $x$-axis along the curve.  The robot pose is chosen based on the largest manipulability measure. The motion primitives are based on equally spaced moveL segments.  We then search for the largest commanded path speed that satisfy both the path uniformity and trajectory tracking constraints as the baseline performance.

\subsection{Simulation setup}
RobotStudio \cite{robotstudio} is an ABB robot dynamics simulator that provides close to real trajectory output. We use a Python interface \cite{abb_motion_program_exec} to RobotStudio virtual controller  to execute motion command directly. In order to achieve high constant speed, the commanded trajectory is extended by first and last command accordingly: MoveL and MoveJ is extended in Cartesian and joint space respectively, and MoveC is extend along the arc and linearly in angle-product orientation. For all motion commands, we start with a blending zone of 10 mm and gradually increase it until the speed profile converges such that the robot does not slow down around waypoints due to blending.

\subsection{Experiment setup}
Fig.~\ref{fig:setup} shows our experiment setup, with the ABB6640 robot holding a mock spray gun. Joint trajectory $\{\mathbf{q}_{exe,i}\}_{i=1}^N$ is recorded through robot controller at 250Hz. We calculate the TCP position and orientation using forward kinematics of the with joint reading. To calculate the position and orientation tracking errrors, for each recorded joint position point $\mathbf{q}_{exe,i}$ we compute
\begin{align}
\mathbf{p}_{err,i} &:=\min_k \norm{\mathbf{p}(\mathbf{q}_{exe,i})-\mathbf{p}_k^*},\\
\mathbf{n}_{err,i} &:= \min_k \arccos \left ( \frac{\mathbf{e}^T_z(\mathbf{q}_{exe,i})\mathbf{n}_k^*}{\norm{\mathbf{e}_z(\mathbf{q}_{exe,i})}\norm{\mathbf{n}_k^*}}\right ).
\end{align}
The maximum position and orientation tracking errors are found by maximizing both terms above over $i$.

\begin{figure}[h]
    \centering
    \includegraphics[width=0.45\textwidth]{images/setup.png}
    \caption{Experiment setup with an ABB6640 robot }
    \label{fig:setup}
\end{figure}

According to  the robot manual, the path repeatability is up to 1.06mm \cite{6640_manual}. To address this issue, for each set of primitive command, we run the robot 5 times and take the interpolated average of 5 recorded trajectories as the execution trajectory.

\subsection{Results}

For baseline execution, in order to satisfy tracking accuracy requirement, we used bisection search on commanded speed until the accuracy requirement is met. For optimized trajectory, we run the post execution optimization to bring the tracking error until it converges. If it does not meet the requirement, we further slow down the commanded speed until the requirement is met.

\begin{table}[h!]
\centering
\begin{tabular}{||c c c c c||}
 \hline
 Curve1 & $max(\mathbf{p}_{err})$ & $max(\mathbf{n}_{err})$ & $\mu(v)$ & $\sigma(v)$ \\ [0.2ex]
 & mm & $^{\circ}$ & mm/s & mm/s \\
 \hline\hline
 Baseline (sim)&            0.49 & 0.18 & 124.01 & 1.02     \\
 Baseline (real)&           0.46 & 0.21 & 103.11 & 0.85     \\
 Bp Adjustment (sim)&       0.16 & 0.57 & 251.29 & 4.77     \\
 Bp Adjustment (real)&      0.38 & 0.56 & 248.77 & 4.33     \\
 Optimized (sim)&           0.16 & 0.35 & 401.23 & 3.55     \\
 Optimized (real)&          0.39 & 0.48 & 395.86 & 8.06     \\[0.5ex]
 \hline
\end{tabular}
\begin{tabular}{||c c c c c||}
 \hline
 Curve2 & $max(\mathbf{p}_{err})$ & $max(\mathbf{n}_{err})$ & $\mu(v)$ & $\sigma(v)$ \\ [0.2ex]
 & mm & $^{\circ}$ & mm/s & mm/s \\
 \hline\hline
 Baseline (sim)&            0.42 & 0.26 & 406.10    & 0.92     \\
 Baseline (real)&           0.47 & 0.28 & 299.97    & 2.36     \\
 Bp Adjustment (sim)&       0.24 & 0.33 & 1089.64   & 51.88    \\
 Bp Adjustment (real)&      0.42 & 0.41 & 973.12    & 13.87    \\
 Optimized (sim)&           0.44 & 0.67 & 1182.36   & 9.62     \\
 Optimized (real)&          0.50 & 1.37 & 1101      & 13.58    \\[0.5ex]
 \hline
\end{tabular}
\caption{Execution Comparison for RobotStudio Simulation and Real Robot Experiments.}
\label{table:results}
\end{table}

From the baseline results, we can see that without any optimization, it is necessary to slow down in order to achieve higher accuracy. Simply adjusting waypoints though post execution optimization could help with increasing traversal speed and minimizing the tracking error. However, we can increase the traversal speed further by optimizing the all redundancies with robot joint velocity and acceleration constraints. The final optimized results show 222\% (sim) 283\% (real) speed increase for Curve 1 and 191\% (sim) and  267\% (real) speed increase for Curve 2 while keeping the tracking accuracy within the requirement.

\subsection{Implementation for a FANUC Robot}

The same algorithm has also been applied to FANUC M710iC-70 robot using the FANUC robot simulation program RoboGuide.
The result is showed in Table \ref{table:results_fanuc}. The speed increased by 323\% for curve1 and 265\% for curve2 while keeping tracking accuracy within the requirement. This shows that the framework successfully generalized to another robot model without fine tuning of the parameters.

\begin{figure}[h]
     \centering
     \begin{subfigure}[b]{0.233\textwidth}
         \centering
         \includegraphics[width=\textwidth]{images/curve_1_robot.png}
         \caption{Curve 1}
     \end{subfigure}
     \begin{subfigure}[b]{0.22\textwidth}
         \centering
         \includegraphics[width=\textwidth]{images/curve_2_robot.png}
         \caption{Curve 2}
     \end{subfigure}
     \caption{FANUC M710iC-70 robot in Roboguide.}
     \label{fig:fanuc_roboguide}
\end{figure}

\begin{table}[h!]
\centering
\begin{tabular}{||c c c c c||}
 \hline
 Curve1 & $max(\mathbf{p}_{err})$ & $max(\mathbf{n}_{err})$ & $\mu(v)$ & $\sigma(v)$ \\ [0.2ex]
 & mm & $^{\circ}$ & mm/s & mm/s \\
 \hline\hline
 Baseline&            0.48 & 1.93 & 98.78 & 1.02     \\
 Bp Adjustment&       0.25 & 1.82 & 319.31 & 7.77     \\[0.5ex]
 \hline
\end{tabular}
\begin{tabular}{||c c c c c||}
 \hline
 Curve2 & $max(\mathbf{p}_{err})$ & $max(\mathbf{n}_{err})$ & $\mu(v)$ & $\sigma(v)$ \\ [0.2ex]
 & mm & $^{\circ}$ & mm/s & mm/s \\
 \hline\hline
 Baseline&            0.43 & 0.92 & 274.61    & 2.48     \\
 Bp Adjustment&       0.38 & 0.75 & 728.45   & 10.83    \\[0.5ex]
 \hline
\end{tabular}
\caption{Execution Comparison in RoboGuide}
\label{table:results_fanuc}
\end{table}

\begin{comment}
\subsection{Generalization Using FANUC Robot}

It is far from enough if the method can only be applied to a specific brand or type of robot arm. In order to test the generalizability of our method, we applied the same framework to FANUC M710iC-70 robot as shown in Fig.~\ref{fig:fanuc_roboguide}.

We performed the experiment in RoboGuide simulation \cite{roboguide}. RoboGuide is the offline programming simulator provided by FANUC. To connect our python code and RoboGuide, a Python interface \cite{fanuc_motion_program_exec} is develop to automate the process of generating FANUC robot TP program and robot execution. Unlike ABB, FANUC defined the blending between motion commands, i.e. positioning path in FANUC's terminology, with two modes, FINE and CNT. The blending is always chosen to be CNT100 to guarantee fastest and smoothest motion.

We applied algorithm described in Sec.~\ref{execution_baseline} for motion program planning. We finally compared the performance between the baseline and waypoint adjustment method discribed in Sec.~\ref{waypoint_adjust}.
The result is showed in Table \ref{table:results_fanuc}. The speed increased by 323\% for curve1 and 265\% for curve2 while keeping tracking accuracy within the requirement. It shows that the framework successfully generalized to another robot model without detail tuning in the algorithm and make generalization to all robot model promising. We hope that in the future we can test the framework on a variety of robot models both in simulation and real robots.

\begin{figure}[h]
     \centering
     \begin{subfigure}[b]{0.233\textwidth}
         \centering
         \includegraphics[width=\textwidth]{images/curve_1_robot.png}
         \caption{Curve 1}
     \end{subfigure}
     \begin{subfigure}[b]{0.22\textwidth}
         \centering
         \includegraphics[width=\textwidth]{images/curve_2_robot.png}
         \caption{Curve 2}
     \end{subfigure}
     \caption{FANUC M710iC-70 robot in Roboguide.}
     \label{fig:fanuc_roboguide}
\end{figure}

\begin{table}[h!]
\centering
\begin{tabular}{||c c c c c||}
 \hline
 Curve1 & $max(\mathbf{p}_{err})$ & $max(\mathbf{n}_{err})$ & $\mu(v)$ & $\sigma(v)$ \\ [0.2ex]
 & mm & $^{\circ}$ & mm/s & mm/s \\
 \hline\hline
 Baseline&            0.48 & 1.93 & 98.78 & 1.02     \\
 Bp Adjustment&       0.25 & 1.82 & 319.31 & 7.77     \\[0.5ex]
 \hline
\end{tabular}
\begin{tabular}{||c c c c c||}
 \hline
 Curve2 & $max(\mathbf{p}_{err})$ & $max(\mathbf{n}_{err})$ & $\mu(v)$ & $\sigma(v)$ \\ [0.2ex]
 & mm & $^{\circ}$ & mm/s & mm/s \\
 \hline\hline
 Baseline&            0.43 & 0.92 & 274.61    & 2.48     \\
 Bp Adjustment&       0.38 & 0.75 & 728.45   & 10.83    \\[0.5ex]
 \hline
\end{tabular}
\caption{Execution Comparison for RoboGuide Simulation Experiments.}
\label{table:results_fanuc}
\end{table}

\subsection{Position Streaming Execution}
The ABB IRB robot series offers an optional External Guided Motion (EGM) interface allowing direct joint position or TCP pose command streaming at 250 Hz \cite{egm_manual}. The tracking accuracy for EGM is not as good as embedded motion primitive commands, and interface may not be universal for robots from other vendors. We have applied iterative learning control (ILC) \cite{deep_rolling} technique for tracking error minimization with baseline and optimized curve pose in RobotStudio simulation, and the result is shown in Table~\ref{table:egm_result}.

\begin{table}[h!]
\centering
\begin{tabular}{||c c c c c||}
 \hline
 Curve1 & $max(\mathbf{p}_{err})$ & $max(\mathbf{n}_{err})$ & $\mu(v)$ & $\sigma(v)$ \\ [0.2ex]
 & mm & $^{\circ}$ & mm/s & mm/s \\
 \hline\hline
 Baseline &            0.49 & 1.82 & 54.96     &  0.26    \\
 ILC &                 0.49 & 2.86 & 230.28    &  1.37     \\
 Optimized &           0.47 & 3.00 & 240.36  & 1.21     \\
 \hline
\end{tabular}
\begin{tabular}{||c c c c c||}
 \hline
 Curve2 & $max(\mathbf{p}_{err})$ & $max(\mathbf{n}_{err})$ & $\mu(v)$ & $\sigma(v)$ \\ [0.2ex]
 & mm & $^{\circ}$ & mm/s & mm/s \\
 \hline\hline
 Baseline &            0.49 & 0.92  & 127.94    & 0.59     \\
 ILC &                 0.50 & 1.63  & 821.81    & 1.28     \\
 Optimized &           0.49 & 2.04  & 902.82   & 0.92     \\
 \hline
\end{tabular}
\caption{EGM Execution Comparison for RobotStudio Simulation and Real Robot Experiments.}
\label{table:egm_result}
\end{table}

The tracking error of EGM pose streaming is much larger compared with motion primitive command, and in order to keep within tracking accuracy, the traversal speed has to slow down to around 50\% of motion primitive baseline. ILC helps with the performance by adjusting streaming pose input, but the performance is not as good as motion primitive program with waypoint adjustment.
\end{comment}

\section{Conclusion and Future Work}
We presented a complete procedure to generate motion primitive commands for industrial robots for tracking high curvature spatial curves with high speed and high accuracy. Our current experiment setup assumes perfect knowledge of the forward kinematics. In the future, we will consider a setup where the TCP position and orientation can be measured with a high precision scanner. Further, we will consider optimization of motion primitives in a dual-arm  setup where the performance is based on the relative TCP frames.

\section*{ACKNOWLEDGMENT}
This research was supported by the ARM Institute project ARM-TEC-21-02-F-19. The authors would like to thank Santiago Paternain and Jonathan Fried for helpful discussion of the project.

\bibliographystyle{IEEEtran}

redpajama_set_name: RedPajamaArXiv


text:
\section{Introduction}

ICD (International Classification of Diseases \cite{icd}) coding refers to the task where medical professionals classify clinical diagnoses and medical procedures associated with each patient using standardised taxonomies, which in turn supports billing, service planning and research. The process is manual and laborious in nature \cite{icderrors}, however there is potential to automate it by identifying relevant information from clinical notes, which are already captured in EHR systems. With this in mind, researchers have started to explore whether machine learning models can succeed at this task \cite{caml}.

The current research on the ICD coding task focuses on the extraction of codes from the discharge summary.
This document is commonly written at the end of a hospital stay and provides a textual description of the important diagnoses and procedures for a given patient, making it particularly helpful for the task.
However, many other clinical notes are also created during the hospital stay, which can provide important details or useful additional context that may be missing from the discharge summary itself.
Analysing the full sequence of notes would allow models to make more accurate decisions and make the problem more similar to a real-life setting, where clinicians have to consider all information about a patient for ICD coding, rather than information only in a single document.

In this work we study how the inclusion of clinical notes across the entire hospital stay can affect performance on the ICD coding task.  We propose the \textbf{H}ierarchical \textbf{T}ransformers for \textbf{D}ocument \textbf{S}equences (\textbf{HTDS}) model, which is an adaptation of the hierarchical transformer model \cite{hibert} for temporal modelling of document sequences. The model takes text and metadata (such as the time and type of note) from a sequence of multiple documents as input and achieves improved performance when additional clinical notes are used for modelling. We compare different prioritisation criteria for selecting which notes to use as input and how to best represent the sequence information. Methods related to superconvergence are applied to speed up the model training process in order to handle the increased size of the data that needs to be processed.

Our experiments show that the inclusion of additional clinical notes indeed improves model accuracy and leads to better predictions.
We evaluate our models against the MIMIC-III-50 \cite{mimic} test set. When considering only the discharge summaries of each hospital stay as input, our model exceeds the current state-of-the-art performance in terms of Micro-F1. When considering all clinical notes as input, further performance improvements across all metrics of interest are observed, exceeding the state-of-the-art performance in Micro-F1, Micro-AUC, Macro-AUC, and Precision@5 scores.

\section{Related Work}

Publicly available electronic health record (EHR) datasets, such as the Medical Information Mart for Intensive Care III (MIMIC-III) dataset \cite{mimic}, provide a shared context for researchers to work on ICD coding. Recent work on ICD coding concentrates on the benchmark tasks presented by \citet{caml}, which extracts ICD codes from the free-text discharge summary generated at the end of each hospital stay. \citet{caml} also publicly release their data preprocessing codes and train/dev/test data splits, and these were followed by later works for comparability of result.

In recent years, state-of-the-art work on the ICD coding problem commonly used methods based on convolutional neural networks (CNNs) or recurrent neural networks (RNNs) for text encoding. CAML \cite{caml} uses a single convolutional layer along with "per-label attention" to extract representations for each label from the convolution output. MSAttKG \cite{msattkg} improves the performance further by using a densely connected convolutional network with variable n-gram features, and incorporating knowledge graphs to capture relationships between medical codes. EffectiveCAN \cite{effectivecan} uses a deep convolutional approach, with a "squeeze-and-excitation" module that repeatedly compresses and then decompresses the convolutional features. LAAT \cite{laat} uses a bidirectional LSTM to encode the texts, with a per-label attention step on the output to get the final classification. MSMN \cite{msmn} uses the same architecture as LAAT, with an additional step of extending code descriptions from the Unified Medical Language System (UMLS) with synonyms, and using an attention layer with a separate head for each code synonym.

Researchers using transformer-based models for text encoding experienced difficulties in matching state-of-the-art performance. \citet{ji} apply a range of different transformer-based models but found that none of them outperformed their reimplementation of a simple CNN-based model. \citet{pascual} similarly found it difficult to achieve competitive performance and concluded that better methods of handling long input sequences are required to improve the models further. \citet{gao} also find that a simple self-attention network with far less parameters outperformed BERT-based models on many tasks. \citet{trldc} show that incorporating task-adaptive pre-training, overlapping chunks, and using a large pretrained language model make it possible to achieve performance that is close to, but still slightly below the state-of-the-art. In general, language models that were pretrained on texts in the biomedicine domain, such as ClinicalBERT \cite{clinicalbert}, BioBERT \cite{biobert}, BlueBERT \cite{bluebert}, and PubMedBERT \cite{pubmedbert} tend to achieve higher performance \cite{trldc, ji} as compared to language models such as BERT \cite{bert} and RoBERTa \cite{roberta} which are trained on general domain corpora, as the models have been adapted to the specialised language used in clinical notes. Among the range of pretrained language models available for the biomedicine domain, better performance was achieved when a specialised token vocabulary is used \cite{pubmedbert,domainroberta} and when the pre-training corpora is closer in nature to those used for the downstream task \cite{dontstoppretraining}. Very recently, \citet{plmicd} identified the restricted capacity of the [CLS] token as a potential limiting factor, and showed how using all tokens in the label attention step leads to state-of-the-art performance on the MIMIC-III-Full problem. However, they do not report results on the MIMIC-III-50 problem.

While transformer-based language models have been very successful on short sequences of text (BERT \cite{bert} and RoBERTa \cite{roberta} use a maximum sequence length of 512 tokens), challenges arise when attempting to apply it to longer text sequences due to the quadratic computational complexity of the self-attention mechanism. Experiments conducted by \citet{gao} show that transformer models require 3x more processing time compared to CNNs, making it more tedious to explore different hyperparameters and modelling strategies. Various modifications have been proposed to the transformer architecture to reduce computation costs, in models such as TransformerXL \cite{transformerxl}, LongFormer \cite{longformer}, and BigBird \cite{bigbird}, however domain-pretrained models for these architectures are relatively scarce. Most transformer-based models for the ICD coding problem adapt the hierarchical transformer \cite{hibert}, which splits the text into chunks that are encoded separately with the pre-trained language model, and then feeds the output of the [CLS] token into a second transformer to allow interaction of information across chunks.

To the best of our knowledge, there has been no prior work that attempts to extend the ICD coding task with other clinical documents.

\section{Approach}

\begin{figure*}[h]
    \centering
    \includegraphics[scale=0.7]{figures/model_architecture_v2.pdf}
    \caption{HTDS Model Architecture. The document sequence is first split into chunks (Step 1) and encoded with a pre-trained language model (Step 2). Meta-information of each chunk is then added to the token encodings (Step 3) before a second transformer is applied to allow attention of information across chunks (Step 4). Finally a label attention layer is applied (Step 5) and the outputs are used for classification (Step 6).}
    \label{fig:architecture}
\end{figure*}

Our \textbf{H}ierarchical \textbf{T}ransformers for \textbf{D}ocument \textbf{S}equences (\textbf{HTDS}) model is based on the hierarchical transformer architecture \cite{hibert}, with additional adaptations specifically to handle document sequences. Figure \ref{fig:architecture}
provides an illustrated diagram of the full HTDS model architecture. We process documents using the following steps:

\textbf{Step 1 - Preprocess and Chunk}: The text in each document is sequentially tokenized and split into chunks, each containing up to $T_c$ tokens. Every new document or clinical note always starts a new chunk.

From these tokenized chunks we select up to $N_c$ chunks for processing. If more than $N_c$ chunks are available, various prioritisation strategies can be considered to select which chunks to use as model input. In our main model we use a strategy that prioritized diversity in the categories of notes used. To do this, we select the last note by timestamp of each category, and then the second last note of each category, and so on until $N_c$ chunks of text are selected.

\textbf{Step 2 - Encode with Language Model}: The chunks are encoded using the pre-trained language model, producing an output of dimension $N_c$ x $T_c$ x $H_e$, where $H_e$ is the dimension of the hidden state in the pre-trained LM.

\textbf{Step 3 - Add Chunk Meta-Information}: Meta-information of each chunk is added. These are learnable embeddings, retrieved via index lookup, with size $H_e$. Positional Embeddings (PE) capture the positional index of the chunk, and are numbered from 0 for the first chunk until N-1 for the last chunk. Temporal Sequence Embeddings (TE) capture the temporal order in which the documents were captured, and are indexed in running order from 0 for chunks belonging to the first document and incremented with each subsequent document. We noted that this indexing approach would often assign varying indices to the last chunk or document, as the number of chunks and documents for each case would vary. This might limit the ability of the model to identify the last chunk or document of the text. Hence, we also include Reversed Positional Embeddings (Rev-PE) and Reversed Temporal Sequence Embeddings (Rev-TE), which start from 0 for the last chunk (or document) and are then incremented with each preceding chunk (or document). Category Embeddings (CE) capture the category of the note, with a unique index for each CATEGORY code. All learnable embeddings use values initialised from a $N(0, 0.1)$ distribution. We hypothesise that these embeddings can help the model to factor in chunk meta-information which may be relevant for classification.

\textbf{Step 4 - Second Transformer}: The embeddings are added together (token embeddings + meta-information embeddings), then concatenated across all the chunks and given as input to a second transformer with $N_e$ encoder layers.
This allows for information from each chunk to interact with all the other chunks and the use of only a small number of layers in this second transformer will keep the computational requirements feasible. The output is an updated embedding of each token, with dimensions ($N_c$ x $T_c$) x $H_e$.

\textbf{Step 5 - Label Attention}: A label attention layer is applied. We train learnable embeddings $\alpha_l$ for each label ($\alpha = [\alpha_1 ... \alpha_{N_l}]$ has dimensions $N_l$ x $H_e$, where $N_l$ is the number of labels) which are applied against the chunk embeddings ($H = [h_1 ... h_{N_c}]$) in an attention step as follows:

\begin{center}
    $$A = softmax(H\alpha^T)$$
    $$V = H^T A$$
    $$Dim(A) = (N_c \times T_c) \times N_l$$
    $$Dim(V) = H_e \times N_l$$
\end{center}

The i-th column in V would be an independent representation, of dimension $H_e$, for the i-th label for classification.

\textbf{Step 6 - Generate Final Classification}: A classification layer is applied. We take $\sigma(W_l v_l)$ to get the probability of label $l$, where $W_l$ is a learnable weight matrix of dimension $H_e$ for label $l$, $v_l$ is the $l$-th item of matrix V, and $\sigma$ is the sigmoid activation function. To obtain the final classification we apply a threshold $t$ for positive classification that is optimised for micro-F1 on the validation set.

\section{Experiment Setup}

\begin{table}[t]
\centering
\begin{tabular}{lcc}
\hline \textbf{} & \textbf{Mean} & \textbf{SD}  \\ \hline
\emph{Discharge Summaries} \\
Total Documents & 1.1 & 0.4 \\
Total Words & 1896 & 929 \\
Total Tokens & 3594 & 1760 \\
{} \\
\emph{All Notes}\\
Total Documents & 33 & 59 \\
Total Words & 10442 & 21334 \\
Total Tokens & 21916 & 46461 \\
\hline
\end{tabular}
\caption{\label{table:summarystats} Summary statistics: Amount of text contained in clinical documents per hospital stay, measured in terms of total number of documents, words, tokens (using the RoBERTa-PM-M3-Voc tokenizer).}
\end{table}

\textbf{Dataset}: For our experiments, we use the MIMIC-III (Medical Information Mart for Intensive Care) dataset \cite{mimic}, which contains multi-modal information on patients admitted to critical care between 2001 and 2012 at the Beth Israel Deaconess Medical Center in Boston, Massachusetts. To limit computational costs, we focus on the MIMIC-III-50 problem, which limits the problem to the top 50 ICD codes by frequency.

To construct the task dataset, we follow \citet{caml} preprocessing steps, with a few exceptions: (1) we keep the text and metadata (specifically the datetime and the category of note) of all notes rather than just the discharge summaries, (2) we do not remove punctuation as we found that performance drops when punctuation is excluded. Each record represents one hospital stay (uniquely identified by the HADM\_ID value) and  contains the texts and ICD codes linked to that hospital stay. There are 8066, 1573 and 1729 records in the train, dev and test sets respectively, giving us a total of 11368 records.

During the data cleaning process, we noticed that the train set contains clinical notes tagged under the category "Nursing/Other", but no clinical notes were tagged in this category in the dev and test sets. For our experiments we grouped "Nursing/Other" and "Nursing" into a single category.

Table \ref{table:summarystats} shows summary statistics of the dataset. In general, discharge summaries are far longer than other documents, with an average of 1724 words per document as compared to the overall average of 316 words per document. However, the text in discharge summaries only accounts for less than 20\% of the words generated in each hospital stay, suggesting the possibility that the other notes might carry additional information that can improve ICD coding accuracy. We also provide the number of tokens produced when the text is tokenized with the RoBERTa-PM-M3-Voc \cite{domainroberta} tokenizer, and we see from the numbers that most hospital stays involve text data that is beyond the 512-token maximum of a single transformer language model.

We also note that the amount of text in each hospital stay can vary widely and has a right-skewed distribution. There is a notable proportion of longer hospital stays which generate substantially more documents and text as compared to the rest. The 90th percentile for Total Words and Total Document Count across all notes is 20556 and 72 respectively. For these hospital stays, the effects of the note prioritisation strategy on model performance would be more prominent.

\textbf{Task Definition}: We investigate two variations of the ICD classification task on this dataset. For Task 1, the notes that are available for modelling are restricted to discharge summaries only. Some hospital stays (11\% of stays) have multiple discharge summaries, typically because of addenda, and in these cases we keep all of them. This would be equivalent to the MIMIC-III-50 task attempted by past works. For Task 2, all notes in each hospital stay are available for use in modelling. This vastly increases the number of documents (from an average of 1.1 to 33 per hospital stay) and the number of words (from an average of 1896 to 10442) to be considered. Task 2 uses the same data splits and labels as Task 1, allowing us to compare the results to assess whether information from the additional notes is able to improve performance.

For both tasks, we use the same evaluation metrics as defined by Mullenbach et al \cite{caml} and then subsequently followed by other researchers: micro-F1, macro-F1, micro-AUC, macro-AUC, and Precision at k=5.

\textbf{Implementation and Model Hyperparameters}: Pytorch was used for the implementation of our models, and NVIDIA Tesla A100 80GB GPUs were used for finetuning. Hyperparameters were tuned manually; Table \ref{table:hyperparameter_search} details the search space and final hyperparameter values used for the HTDS model. The pretrained language model was initialised to the RoBERTa-base-PM-M3-Voc \cite{domainroberta} model checkpoint, which was pretrained on texts in PubMed, PubMed Central, and MIMIC-III physician notes. The second transformer uses 1 encoder layer with 8 attention heads.

\begin{table}[t]
\begin{tabular}
{p{3.3cm}p{3.3cm}}
\hline  Hyperparameter & Values\\
\hline
\emph{Optimization} \\
Peak Learning Rate & 1e-6 to 1e-4 \textbf{(5e-5)}\\
Number of Epochs & 10-50 \textbf{(20)} \\
Early Stopping Patience Threshold & None, \textbf{3}, 5, 10 \\
Effective Batch Size & 1-64 \textbf{(16)} \\
\\
\emph{Language Model} \\
Pre-trained LM & PubMedBERT, \textbf{RoBERTa-base-PM-M3-Voc}, RoBERTa-large-PM-M3-Voc\\
Tokens per chunk, $T_c$ & \textbf{512}\\
Max Chunks, $N_c$ & 1-48 \textbf{(32)}\\
\\
\emph{Second Transformer} \\
Encoder Layers & 0, \textbf{1}, 2 \\
Attention Heads & \textbf{8}, 12 \\
\hline
\end{tabular}
\caption{\label{table:hyperparameter_search} Hyperparameter search space. Bolded text indicates hyperparameters used in the HTDS model.}
\end{table}

Texts are tokenized into chunks of $T_c$=512 tokens and a maximum of $N_c$=32 chunks were used as model input. With these values for $T_c$ and $N_c$, the note selection strategy to maximise diversity of document categories (detailed earlier in Section 3) was applied for 45\% of samples which have more than 32 chunks of text. The model has 136M parameters in total.

\begin{table*}[h]
\begin{center}
\begin{tabular}
{lccccc}
\hline  {} & Micro $F_1$ & Macro $F_1$ & Micro AUC & Macro AUC & P@5 \\
\hline
\multicolumn{2}{l}{\emph{CNN-based Models}} \\
CAML \cite{caml} & 63.3 & 57.6 & 91.6 & 88.4 & 61.8 \\
MSAttKG \cite{msattkg} & 68.4 & 63.8 & 93.6 & 91.4 & 64.4 \\
EffectiveCAN \cite{effectivecan} & 71.7 & 66.8 & 94.5 & 92.0 & 66.4 \\
\hline
\multicolumn{2}{l}{\emph{RNN-based Models}} \\
LAAT \cite{caml} & 71.5 & 66.6 & 94.6 & 92.5 & 67.5 \\
MSMN \cite{msmn} & 72.5 & \textbf{68.3} & 94.7 & 92.8 & 68.0 \\
\hline
\multicolumn{2}{l}{\emph{Transformer-based Models}} \\
Hier-PubMedBERT \cite{ji} & 68.1 & 63.3 & 90.8 & 88.6 & 64.4 \\
TrLDC (Base) \cite{trldc} & 70.1 & 63.8 & 93.7 & 91.4 & 65.9 \\
TrLDC (Large) \cite{trldc} & 71.1 & 65.5 & 94.1 & 91.9 & 66.4 \\
\hline
\multicolumn{2}{l}{\emph{Our Models}} \\
HTDS (Discharge Summaries) & $72.6_{0.3}$ & $66.6_{1.2}$ & $94.5_{0.1}$ & $92.6_{0.3}$ & $67.4_{0.3}$ \\
HTDS (All Notes) & \textbf{73.3}$_{0.3}$ & $67.9_{0.4}$ & \textbf{95.0}$_{0.2}$ & \textbf{93.2}$_{0.2}$ & \textbf{68.2}$_{0.2}$ \\
\hline
\end{tabular}
\caption{\label{table:main_results} Performance of models on the MIMIC-III-50 test set. Models are sorted by Micro-F1 within each category. Metrics are averaged across 5 replications. Subscripts indicate the standard deviation across runs. Bolded values indicate the best score achieved for each metric.}
\end{center}
\end{table*}

These hyperparameters were selected to maximise Micro-F1 on the dev set, with a few exceptions to manage training and computation costs: (1) while using the larger RoBERTa-large-PM-M3-Voc model was found to achieve better performance, we kept to the smaller RoBERTa-base-PM-M3-Voc model; (2) while increasing the maximum number of chunks $N_c$ in general leads to better performance, we limit our model to a maximum of 32 chunks.

Training models that take text across all clinical documents as inputs, compared to using only the discharge summary, requires substantially more computational resources. With A100 GPUs, 15.5 samples are processed per second when training on discharge summaries only\footnote{TrLDC \cite{trldc}, which we consider to be a comparable model in terms of architecture, processed 7.4 samples per second when training on discharge summaries on NVIDIA V100 GPUs.}, and 4.9 samples are processed per second when training with all clinical documents.
To speed up the model optimisation process, we apply the 3-phase \emph{1cycle}
learning rate scheduler for superconvergence as described in \cite{superconvergence}. The learning rate (LR) progresses via cosine annealing from 1/25 of the peak LR to the peak LR (5e-5) in the first phase (30\% of iterations) and then goes back to 1/25 of the peak LR in the second phase (30\% of iterations). Finally in the third phase (40\% of iterations), LR is annealed to 1/1000 of the peak LR. The AdamW optimizer is used, with an effective batch size of 16 achieved through gradient accumulation. The model is trained for up to 20 epochs with an early stopping patience threshold of 5. With this setup, training is stopped at around the 14th epoch on average. We note that this is at least 50\% less (in terms of number of epochs) compared to past works on the MIMIC-III-50 problem where transformer-based models would be trained for 30 epochs or more \cite{trldc, ji, pascual}.

\section{Results}

\subsection{Main Results} Table \ref{table:main_results} shows the results when our models are evaluated against the MIMIC-III-50 test set, as well as comparisons against published works. We report the averaged metrics across 5 training replications.

As we can see from the table, prior works with transformer-based models faced challenges in achieving competitive performance on this problem. \citet{trldc} managed substantial improvements with the TrLDC model over the work of \citet{ji}, however even with a large-sized model their performance still fell slightly behind the best-performing CNN-based and RNN-based models. When using only discharge summaries, HTDS achieves state-of-the-art performance in terms of Micro-F1, the primary metric used for comparison. It also exceeds past CNN-based and Transformer-based models on all metrics of interest.

When including all clinical documents, as compared to including only discharge summaries, the performance of HTDS improves on all metrics of interest (all differences are statistically significant at p<0.05), including an additional 0.7\% increase in Micro-F1. Comparing against all other models, we see that the model achieves state-of-the-art performance in terms of all metrics except for Macro-F1. We hypothesize that the modelling of code synonyms in MSMN \cite{msmn} helped to increase its performance on rarer ICD codes and hence achieve a higher Macro-F1 score, but also note that steps used to improve performance by incorporating synonyms based on UMLS concepts could also be adapted into our model to achieve similar improvements.

Put together, our results demonstrate the value of including clinical documents beyond the discharge summary in modelling.

\subsection{Ablation Experiments}
To analyse the effect of various components and hyperparameter choices on model performance, we start with our main model and then ablate or vary individual components one at a time, keeping all other components constant, and evaluate their performance on the development set. We share our results in this section.

For all ablation experiments, we report the impact on Micro-F1, the primary metric of interest, averaged across 5 replications.

\textbf{Quantity of Text Input}:
Table \ref{table:qty_input} shows how performance varies as the quantity of text is varied. The quantity of text used as input has a substantial impact on the compute requirements of the entire model architecture. When $N_c$ is reduced 16, 7.5 samples are processed per second when training on A100 GPUs, an increase of 0.5x as compared to 4.9 samples per second for HTDS which uses $N_c$=32. However, as we can see from the results of this ablation experiment, reducing the quantity of text input leads to a substantial drop in model performance.

\begin{table}[h]
\begin{tabular}
{p{5.1cm}cc}
\hline  {} & Micro $F_1$\\
\hline
HTDS (Max 32 Chunks) & 74.0\\
Max 16 Chunks & 73.0\\
\hline
\end{tabular}
\caption{\label{table:qty_input} Performance when the quantity of text input is varied on the development set.}
\end{table}

\textbf{Metadata embeddings}: Table \ref{table:metadata embeddings} shows how the performance varies as the metadata embeddings used in the model are varied. The ablation of each of the embedding types in isolation results in small but consistent decreases in model performance. It is possible that the model compensates by learning to capture some of this information from the text itself without explicit embeddings. Indeed, past works have observed that the clinical notes in the MIMIC-III dataset have a high degree of structure and templating \cite{ehrwriting}. Nevertheless, in our experiments the overall best results were obtained by using the combination of all the proposed metadata embeddings.

\begin{table}[h]
\begin{center}
\begin{tabular}
{p{5.1cm}ccccc}
\hline  {} & Micro $F_1$\\
\hline
HTDS (All meta embeddings) & 74.0 \\
Ablate CE & 73.9 \\
Ablate PE+Rev-PE & 73.9\\
Ablate TE+Rev-TE & 73.8\\
\hline
\end{tabular}
\caption{\label{table:metadata embeddings} Performance when metadata embeddings are varied on the development set. }
\end{center}
\end{table}

\textbf{Chunk Representations}: In a traditional hierarchical transformer, only the encoding of the [CLS] token is kept and used as an aggregate representation of the chunk. However, recent works have suggested that the [CLS] token might have insufficient capacity to capture information for the large number of labels in the ICD coding problem \cite{plmicd}. In Table \ref{table:chunk_representation}, we show the results when only the [CLS] token is used as an aggregate representation of each chunk, and see that there is a sizeable decrease in performance.

\begin{table}[h]
\begin{center}
\begin{tabular}
{p{5.1cm}ccccc}
\hline  {} & Micro $F_1$\\
\hline
HTDS (All token representations) & 74.0 \\
CLS token representation only & 71.7 \\
\hline
\end{tabular}
\caption{\label{table:chunk_representation} Performance when the embeddings used for chunk representation are varied on the development set. }
\end{center}
\end{table}

\textbf{Second Transformer}: The second transformer in Step 4 allows tokens from each chunk to attend to tokens from other chunks. While earlier studies \cite{trldc,ji} include this second transformer, it also adds to the computational costs of the model due to the quadratic complexity of the attention step and \cite{plmicd} show that the second transformer can be dropped if the encodings of all tokens (rather than just the [CLS] token) are kept for the label attention step.

Our ablation experiments in Table \ref{table:second_transformer} provide some additional insight on this. When considering only the discharge summary, the second transformer can be dropped without substantial impact on performance. However, when modelling the sequence of all clinical documents, ablating the second transformer leads to a noticeable decrease in performance, suggesting that the information in other documents can help further refine token representations before classification.

\begin{table}[h]
\begin{center}
\begin{tabular}
{p{5.1cm}ccccc}
\hline  {} & Micro $F_1$ \\
\hline
HTDS (Discharge Summaries) & 73.2 \\
Ablate 2nd Transformer & 73.3 \\
\\
HTDS (All Notes) & 74.0 \\
Ablate 2nd Transformer & 73.6 \\
\hline
\end{tabular}
\caption{\label{table:second_transformer} Performance when second transformer is ablated on the development set.}
\end{center}
\end{table}

\textbf{Note Selection:}
In around 45\% of admissions, tokenizing the text in all available clinical notes will produce more than 32 chunks of 512 tokens. In those cases, we would need to select which chunks are used as inputs to our model. Table \ref{table:note selection} shows our results. We considered the following strategies to prioritise which chunks to use:
\begin{itemize}
    \item By timestamp: We select the first or last 32 chunks by the timestamp of the clinical notes. Taking the last chunks achieved far superior performance.
    \item By category: We select first the discharge summary\footnote{Our exploratory tests find that the discharge summaries contain the most relevant information. We note also that prior work achieved good performance with just the discharge summaries, without the need for other notes.}, then notes of a certain category (Radiology/Nursing/Physician), and then other notes until 32 chunks of text are selected. Our results indicate that the differences in performance are mostly marginal, suggesting that there could be multiple possible strategies that achieve close to optimal performance.
    \item Prioritise diversity:  We select first the last note by timestamp of each category, and then the second last note of each category, and so on until 32 chunks of text are selected. This maximises the diversity (in terms of categories of notes) used as inputs to the model. This approach was found to have the highest score on the development set, and hence used for HTDS.
\end{itemize}

\begin{table}[h]
\begin{center}
\begin{tabular}
{p{5.1cm}cc}
\hline  {} & Micro $F_1$ \\
\hline
HTDS (Prioritise diversity) & 74.0 \\
\\
Prioritise First & 68.4\\
Prioritise Last & 73.8 \\
\\
Prioritise Radiology & 73.8 \\
Prioritise Nursing & 73.7 \\
Prioritise Physician & 73.8 \\
\hline
\end{tabular}
\caption{\label{table:note selection} Performance when note selection is varied on the development set. }
\end{center}
\end{table}

In general, we also note that the effects of note selection strategies would be more pronounced when the maximum number of chunks $N_c$ for model input is smaller, as it would result in a greater proportion of text being excluded.

\section{Conclusion}

As we work towards automated ICD coding, it would be helpful to build models that can consider information that is captured across the patient's EHR record, rather than just the discharge summary (which may not always be exhaustive). Such an approach would also be more similar to a real-life setting, where clinicians consider all available information for ICD coding, rather than information in a single document.

In this paper, we demonstrated the HTDS model, which is an adaptation of the hierarchical transformer model that considers the text and metadata from the entire clinical document sequence for ICD coding. While transformer-based models have faced difficulties achieving competitive performance on the ICD coding problem in the past, with HTDS we show that these challenges can be overcome. When evaluated on the MIMIC-III-50 test set using only discharge summaries, HTDS exceeded the prior state-of-the-art performance in terms of Micro-F1 score. When all clinical documents were considered, the performance of HTDS improved further on all performance metrics of interest, and exceeded prior state-of-the-art performance in all metrics except for Macro-F1. The results demonstrate the value of including clinical documents beyond the discharge summary in the ICD coding problem.

Possibilities for improving performance even further are plenty. These include: using a large-sized language model or using overlapping text chunks to reduce fragmentation in the initial encoding step \cite{trldc}, considering other transformer architectures for long texts \cite{longformer, transformerxl, bigbird}, smarter strategies for chunking the input text to reduce fragmentation, further improving the strategy for selecting which text to use as model input (possibly going down to text-level rather than document-level approaches), and incorporating methods to better model rare ICD codes \cite{laat, msmn}. Approaches for improving the computational efficiency and training time of the model can be considered to help to reduce GPU resource requirements, and enable the testing of more models and hyperparameter settings. Going even further from here, we could consider multi-modal models that use information across the entire EHR database for ICD coding.

We hope that our findings will encourage future studies that tap on the full breadth and depth of information available in modern EHR databases today in order to further push the limits of performance on the ICD coding problem in future.

\section*{Limitations}
Although applying HTDS on the full clinical document sequence in each hospital stay helped to push performance on the ICD coding problem further as compared to the prior state-of-the-art, we note a few limitations to our work.

Firstly, the computational requirements to train HTDS is not trivial. When using NVIDIA A100 GPUs, one training run took 8 GPU-hours on average (for 5 replications this would require 40 GPU-hours). The increased computation cost for HTDS, as compared to other models on the ICD coding problem, could be attributed to the higher number of model parameters in transformers as compared to CNN/RNNs and the increase in input data size as a result of using all clinical documents as input. It is hoped that this issue of high compute costs can be mitigated in future by either further refinements in modelling to improve efficiency or improvements in the compute capabilities of hardware used for model training.

Secondly, we note that our work focuses only on the MIMIC-III-50 problem, where only the top 50 ICD codes by frequency are considered. This would be insufficient in a real-life setting, which would require clinicians to consider all ICD codes. Extending our work on the MIMIC-III-Full problem, which uses a dataset that is 4x in size, was not attempted due to limitations on compute resources. However, we speculate that the benefits of using all clinical documents to perform ICD coding would apply similarly to the MIMIC-III-Full problem.

Finally, while we have taken the actual ICD codes assigned by clinicians as the "ground truth" for the purpose of model evaluation, there have been errors made during the process. We would not expect clinicians to thoroughly read the entire clinical document sequence (consisting an average of over 10,000 words) for every patient when performing ICD coding, and hence there is a possibility that some codes could have been missed. A more thorough approach for model evaluation could involve extracting a sample of records where different codes were assigned by the clinicians and our models for further evaluation by experts, in order to determine the extent to which this might have affected our evaluation metrics.

\section*{Ethics Statement}
No conflicts of interest are declared by the authors. Clinical data in the MIMIC-III database is de-identified through removal of identifying data elements and date-shifting in accordance with Health Insurance Portability and Accountability Act (HIPAA) standards, and protected health information was further removed from clinical notes using dictionary look-ups and pattern-matching \cite{mimic}. The use of the data is in accordance with the MIMIC-III data use agreement.


text:
\section{Introduction}

The structure of the proton is a result of complicated non-perturbative many-body interactions between its fundamental building blocks, the quarks and gluons. This structure is encoded in various distribution functions, the simplest one being the collinear parton distribution functions that describe the parton density as a function of longitudinal momentum fraction $x$ carried by the parton (measured at a given scale $Q^2$). These distributions have been measured with great precision by experiments at HERA~\cite{Aaron:2009aa} by studying total electron-proton cross sections.

More detailed information can be extracted from more differential observables that provide access to more differential distribution functions. For example, in exclusive photon or vector meson production the total momentum transfer is measurable, and via Fourier transform provides access to the spatial distribution of partons known as the generalized parton distribution function GPDF~\cite{Diehl:2003ny,Belitsky:2005qn}. It is also possible to study the distribution of partons in the proton in transverse momentum space, described by transverse momentum dependent parton distribution functions (TMDs)~\cite{Collins:1981uw}.

The most complete information of the proton structure is encoded in the Wigner distribution~\cite{Ji:2003ak,Belitsky:2003nz,Lorce:2011kd}, which depends on both transverse coordinate and transverse momentum, in addition to the longitudinal momentum fraction $x$. This quantum distribution is not positive definite and has a probabilistic interpretation only in certain semi-classical limits \cite{Moyal:1949sk,Hillery:1983ms,Polkovnikov:2009ys}.
To access the Wigner distribution, more differential observables than single particle production or total cross sections are needed. In Ref.~\cite{Hatta:2016dxp}, it was shown that diffractive dijet production, where two jets are produced in a process where no net color charge is exchanged with the target, is sensitive to the gluon Wigner distribution at small $x$. A future Electron Ion Collider in the US~\cite{Accardi:2012qut,Aschenauer:2017jsk} or LHeC~\cite{AbelleiraFernandez:2012cc} at CERN would be able to measure this process over a wide kinematical region at high center-of-mass energies.

At high energies or small $x$ the convenient effective theory to describe high energy scattering processes is provided by the Color Glass Condensate (CGC), which describes Quantum Chromodynamics (QCD) in the high energy limit. In Ref.~\cite{Mantysaari:2019csc}, summarized here, we calculate both the diffractive dijet production cross section and the Wigner distribution in the CGC framework.

\section{Dipole-proton interaction in the CGC}
\label{sec:dipole}
At high energies the convenient degrees of freedom are Wilson lines $U(\mathbf{x})$ that describe the color rotation that the parton encounters when propagating eikonally through the target. For a given target configuration, the Wilson lines are obtained by solving the Yang-Mills equations
\begin{equation}
U(\mathbf{x}) = P \exp \left( -ig \int \mathrm{d} x^- \frac{\rho (x^-,\mathbf{x})}{\nabla^2 + \tilde m^2 }\right).
\end{equation}
The color charge density $\rho$ is assumed to be a local Gaussian variable with expectation value being related to the local density of the proton, which we assume to be Gaussian in this work, and $\tilde m^2$ is an infrared regulator. After the Wilson lines are determined at the initial Bjorken-$x$, their evolution to smaller $x$ is obtained by solving the perturbative JIMWLK evolution equations (see e.g.~\cite{JalilianMarian:1996xn}).
  All parameters that control e.g. the density at the initial $x_0=0.01$, the size of the proton and the values of the strong coupling and infrared regulators are constrained by the HERA structure function and diffractive vector meson production measurements~\cite{Mantysaari:2018zdd}.
  For a more detailed description of the setup, the reader is referred to~\cite{Mantysaari:2019csc}.

When Wilson lines are sampled on the lattice and evolved to smaller $x$ with the JIMWLK equation, it becomes possible to construct the dipole-target scattering amplitude at any $x$
\begin{equation}
N\left( \mathbf{r} = \mathbf{x} - \mathbf{y}, \mathbf{b} = \frac{\mathbf{x} + \mathbf{y}}{2} \right) = 1 - \frac{1}{N_c} \langle \mathrm{Tr} U(\mathbf{x}) U^\dagger(\mathbf{y}) \rangle,
\end{equation}
where the average is taken over different possible target configurations.

When we consider exclusive dijet production, the Fourier conjugates to the dijet momentum and to the recoil momentum are the dipole size $\mathbf{r}$ and  impact parameter $\mathbf{b}$. Having this in mind, we study the angular modulation of the dipole-proton scattering amplitude $N(\mathbf{r},\mathbf{b})$ calculated from the CGC framework. The dipole amplitude $N$ as a function of the angle between $\mathbf{r}$ and $\mathbf{b}$ is shown in Fig.~\ref{fig:dipole}. Note that in widely used dipole amplitude parametrizations such as IPsat~\cite{Kowalski:2003hm} there would be no angular dependence.

To quantify the evolution of the elliptic modulation of the dipole amplitude, we extract the Fourier harmonics $v_n$ writing the dipole amplitude as
$	N(\mathbf{r},\mathbf{b}) = v_0 [1 + 2 v_2 \cos 2\theta(\mathbf{r},\mathbf{b}) ]$.
The extracted $v_0$ and $v_2$ coefficients at different rapidities are shown in Fig.~\ref{fig:dipole}. We find that the evolution suppresses the elliptic modulation (note that Bjorken-$x$ is related to the evolution rapidity as $x = x_0 e^{-y}$ with $x_0=0.01$). This is mainly due to the rapid growth of the proton density in the dilute region, resulting in a smoother and larger proton with smaller density gradients

 \begin{figure}[tb]
    \centering
    \begin{minipage}{.48\textwidth}
        \centering
        \includegraphics[width=\textwidth]{dipole_theta.pdf}
				\caption{Dipole amplitude as a function of angle between dipole size and impact parameter at different rapidities. Figure from Ref.~\cite{Mantysaari:2019csc}.}

		\label{fig:dipole}
    \end{minipage} \quad
    \begin{minipage}{0.48\textwidth}
        \centering
      \includegraphics[width=\textwidth]{husimi_Pdep_evolution}
				\caption{Rapidity evolution of the elliptic component of the Husimi distribution from Ref.~\cite{Mantysaari:2019csc}.}
		\label{fig:husimi}
    \end{minipage}
\end{figure}

\section{Wigner and Husimi distributions}
As discussed in the Introduction, the gluon Wigner distribution $xW(\mathbf{P}, \mathbf{b}, x)$ contains the most  complete information of the small-$x$ gluonic structure of the proton. In particular, it describes the gluon distribution as a function of both transverse coordinate $\mathbf{b}$ and transverse momentum $\mathbf{P}$. The disadvantage is that due to the uncertainty principle it can not have a probabilistic interpretation, and we indeed show in Ref.~\cite{Mantysaari:2019csc} that when calculated from the CGC framework, the Wigner distribution becomes negative at small transverse momenta.

If the Wigner distribution is smeared over both transverse coordinate and transverse momentum with the smearing parameters being inverse to each other, one obtains the so called Husimi distribution
\begin{equation}
xH(\mathbf{P},\mathbf{b},x)) = \frac{1}{\pi^2} \int \mathrm{d}^2 \mathbf{b}' \mathrm{d}^2 \mathbf{P}' e^{-\frac{1}{l^2}(\mathbf{b}-\mathbf{b}')^2 - l^2(\mathbf{P}-\mathbf{P}')^2} xW(\mathbf{P}',\mathbf{b}',x).
\end{equation}
In this work we choose $l=1\,\mathrm{GeV}^{-1}$, as it corresponds to a distance scale much smaller than the proton size, but does not result in too large smearing in momentum space that would wash out most of the transverse momentum dependence. In Ref.~\cite{Mantysaari:2019csc} it is shown that the Husimi and Wigner distributions agree at large $|\mathbf{P}| \gtrsim 1/l$, and that the Husimi distribution calculated from the CGC framework following~\cite{Hagiwara:2016kam} is positive definite.

To study the elliptic modulation (dependence on the angle between $\mathbf{P}$ and $\mathbf{b}$), we write the Husimi distribution as
\begin{equation}
xH(\mathbf{P},\mathbf{b},x) = v_0^H [1 + 2 v_2^H \cos 2\theta(\mathbf{P},\mathbf{b})].
\end{equation}
The rapidity evolution of the elliptic coefficient $v_2^H$  is shown in Fig.~\ref{fig:husimi}. Except at the smallest momenta, the evolution suppresses the elliptic modulation as expected based on the analysis of the dipole amplitude in Sec.~\ref{sec:dipole}. At the smallest momentum values the elliptic component first grows in the evolution, which can be understood as when the proton grows, large dipoles $|\mathbf{r}| \sim |\mathbf{P}|^{-1}$ with large elliptic modulation start to contribute, before the proton grows enough that we start to see the decreasing density gradients also at small $|\mathbf{P}|$.

\section{Diffractive dijet production}
As discussed in Ref.~\cite{Hatta:2016dxp}, diffractive dijet production is sensitive to the gluon Wigner distribution. In particular, it is interesting to study diffractive dijet production as a function of the two momentum vectors, the average momentum $\mathbf{P} = \frac{1}{2}(\mathbf{p}_1 - \mathbf{p}_2)$ and the recoil momentum $\boldsymbol{\Delta} = \mathbf{p}_1 + \mathbf{p}_2$, where $\mathbf{p}_1$ and $\mathbf{p}_2$ are the momenta of the individual jets. The connection to the Wigner distribution is shown in the correlation limit $|\mathbf{P}| \gg |\boldsymbol{\Delta}|$, thus we use $|\boldsymbol{\Delta}|=0.1$ GeV in this work.

The diffractive dijet production cross section in the CGC framework is derived in Ref.~\cite{Altinoluk:2015dpi}. The cross section as a function of jet momentum $\mathbf{P}$ is shown in Fig.~\ref{fig:dijet_spectrum}, where results obtained with different infrared regulators (with value of the coupling constant adjusted to describe the HERA structure function data) and using both fixed and running coupling evolution are shown, and our results are found to be insensitive to the infrared regularization. Here it is worth noticing that the Fourier conjugate to $\mathbf{P}$ is the dipole size, and thus this can be seen as a diffraction off the $q\bar q$ Fock state of the probing photon. Here we only consider charmed dijets, so the diffractive dip location can be estimated to be $|\mathbf{r}_\gamma|^{-1} \sim \sqrt{m_c^2 + Q^2} \approx 1.4\,\mathrm{GeV}$.

To study the elliptic modulation we extract the Fourier harmonics of the dijet production cross section:
\begin{equation}
\mathrm{d}\sigma = v_0[1 + 2 v_2 \cos 2\theta(\mathbf{P},\boldsymbol{\Delta})].
\end{equation}
The elliptic coefficient $v_2$ as a function of Bjorken-$x$ (denoted as $x_p$)  is shown in Fig.~\ref{fig:dijet_v2}, where we find that the energy evolution reduces $v_2$ by almost a factor of $2$ in the EIC energy range. This is mostly due to the increasing proton size suppressing density gradients. The modulation is relatively small and likely difficult to measure. However, at larger $|\boldsymbol{\Delta}|$ we expect a much larger signal~\cite{Salazar:2019ncp}. For comparison, we also show the result obtained in the case where we do not perform the JIMWLK evolution towards small-$x$, but just scale the overall proton density, in which case the $v_2$ is independent of energy. In the models where there is no dependence on the angle between $\mathbf{r}$ and $\mathbf{b}$ in the dipole amplitude, one gets exactly $v_2=0$~\cite{Altinoluk:2015dpi}.

 \begin{figure}[tb]
    \centering
    \begin{minipage}{.48\textwidth}
        \centering
        \includegraphics[width=\textwidth]{PdepOverview_IPGlasma_charm_longitudinal.pdf}
				\caption{Charmed dijet photoproduction spectrum as a function of dijet momentum. Figure from Ref.~\cite{Mantysaari:2019csc}.}
		\label{fig:dijet_spectrum}
    \end{minipage} \quad
    \begin{minipage}{0.48\textwidth}
        \centering
      \includegraphics[width=\textwidth]{xdep.pdf}
				\caption{Elliptic ($v_2$) modulation of dijet photoproduction as a function of momentum fraction $x_p$ of the target. The values on top refer to the  center of mass energies $W$. Figure from Ref.~\cite{Mantysaari:2019csc}.}
		\label{fig:dijet_v2}
    \end{minipage}
\end{figure}

\section{Conclusions}

We have calculated the Wigner and Husimi distributions from the CGC framework and the diffractive dijet production cross section, which in principle is sensitive to the gluon Wigner distribution at small $x$.
By solving the perturbative JIMWLK evolution equations we find that in the gluon distributions the correlation between the momentum and coordinate space angles decreases when evolving towards small-$x$, and predict that the elliptic modulation in the dijet production cross section decreases almost by a factor of $2$ from the lowest to highest center of mass  energies at the EIC.

\subsection*{Acknowledgements}
HM is supported by the Academy of Finland, project 314764, and by the European Research Council, Grant ERC-2015-CoG-681707. NM and BS are supported by the U.S. Department of Energy, Office of Science, Office of Nuclear Physics, under contract No.DE- SC0012704. NM is funded by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) - Project 404640738.

\bibliographystyle{JHEP}