Optimizing Drug Delivery to Tumors

A Two-Compartment Pharmacokinetic Approach

Why particle size and infusion rate matter

Nanoparticle-based cancer therapeutics rely on the Enhanced Permeability and Retention (EPR) effect: tumor blood vessels are leaky and poorly drained, so nanoparticles of the “right” size accumulate preferentially inside tumor tissue. But size cuts both ways. Particles that are too small are filtered out by the kidneys before they can accumulate anywhere. Particles that are too large are swept up by the reticuloendothelial system (liver, spleen) and cleared from circulation quickly. There is a sweet spot, typically somewhere in the 80–150 nm range, where a particle circulates long enough and permeates the tumor efficiently enough to deliver a meaningful dose.

At the same time, the infusion rate is not a free parameter — push too much drug into the blood and you hit systemic toxicity limits before the tumor ever gets a useful concentration.

This post builds a small, fully quantitative model that captures both effects, then numerically finds the particle size and infusion rate that maximize tumor drug exposure without violating a plasma toxicity constraint.

The model

We treat the body as two compartments: plasma and tumor tissue. Drug crosses between them at a rate set by the permeability–surface area product $PS$, which depends on particle diameter $d$ through the EPR effect:

$$PS(d) = PS_{max}\exp!\left(-\frac{(d-d_{opt})^2}{\sigma_d^2}\right)$$

Clearance from plasma also depends on particle size — small particles are cleared fast by renal filtration, large ones by the RES, and clearance is minimized near some optimal circulation diameter $d_c$:

$$CL(d) = CL_0\left(1+\left(\frac{d-d_c}{w_c}\right)^2\right)$$

With these two size-dependent parameters, the plasma concentration $C_p$ and tumor concentration $C_t$ evolve as:

$$\frac{dC_p}{dt} = \frac{R_{in}}{V_p} - \frac{CL(d)+PS(d)}{V_p}C_p + \frac{PS(d)}{V_p}C_t$$

$$\frac{dC_t}{dt} = \frac{PS(d)}{V_t}\left(C_p-C_t\right) - k_{deg}C_t$$

where $R_{in}$ is the (zero-order) infusion rate, $V_p$ and $V_t$ are the plasma and tumor volumes, and $k_{deg}$ is the tumor drug elimination rate.

This is a linear system $\dot{\mathbf{x}} = A\mathbf{x} + \mathbf{b}$, so it has a closed-form solution — no numerical ODE integration needed at all. We use this to evaluate the whole model for a full grid of particle sizes at once, using nothing but NumPy array broadcasting.

The metric we optimize is the tumor drug exposure, i.e. the area under the tumor concentration curve:

$$AUC_{tumor} = \int_0^{T_{end}} C_t(t),dt$$

subject to a plasma toxicity ceiling:

$$\max_{d,,R_{in}}\ AUC_{tumor}(d,R_{in}) \quad \text{s.t.} \quad \max_t C_p(t) \le C_{tox}$$

Why we can avoid a slow parameter sweep

A naive implementation would loop over every (particle size, infusion rate) pair — say 300 × 200 = 60,000 combinations — and call scipy.integrate.solve_ivp for each one. That’s 60,000 separate numerical integrations, which is slow and wasteful.

We avoid this in two ways:

  1. Analytical matrix exponential. For a 2×2 system, $e^{At}$ has a closed form in terms of the eigenvalues of $A$ (Sylvester’s formula), so we never call a generic matrix-exponential routine — we compute it directly with array operations, broadcast simultaneously over every particle size and every time point.
  2. Linearity in the infusion rate. Because the system is linear, the entire trajectory scales linearly with $R_{in}$. We only need to solve the model once per particle size at a unit infusion rate; every other infusion rate is just a multiplication. This collapses a genuinely 2D sweep into a 1D computation plus simple algebra.

The result: the full parameter sweep over 300 particle sizes and 200 infusion rates runs in a fraction of a second, with zero explicit loops over the parameter grid.

Source code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
import numpy as np
import matplotlib.pyplot as plt

# ---------------------------------------------------------------
# 1. Physiological / model parameters (illustrative values)
# ---------------------------------------------------------------
CL0 = 0.5 # baseline plasma clearance at optimal circulation size [L/h]
d_c = 130.0 # particle diameter with longest circulation time [nm]
w_c = 60.0 # width of the clearance "bowl" [nm]

PS_max = 0.02 # max permeability-surface-area product (EPR effect) [L/h]
d_opt = 100.0 # particle diameter with best tumor permeability [nm]
sigma_d= 40.0 # width of the EPR window [nm]

Vp = 5.0 # plasma volume [L]
Vt = 0.01 # tumor volume [L] (10 mL tumor)
kdeg = 0.05 # tumor drug elimination rate [1/h]

Tinf = 2.0 # infusion duration [h]
Tend = 48.0 # total observation window [h]
Nt = 2000 # number of time points
C_tox = 10.0 # plasma toxicity threshold [mg/L]

t = np.linspace(0, Tend, Nt)
t1_mask = t <= Tinf
t1 = t[t1_mask]
t2 = t[~t1_mask] - Tinf # time measured from the end of infusion

# ---------------------------------------------------------------
# 2. Size-dependent physiology
# ---------------------------------------------------------------
def PS_of_d(d):
"""EPR-driven permeability-surface-area product vs particle size."""
return PS_max * np.exp(-((d - d_opt) / sigma_d) ** 2)

def CL_of_d(d):
"""Plasma clearance vs particle size (renal loss when small,
RES uptake when large; minimized near d_c)."""
return CL0 * (1 + ((d - d_c) / w_c) ** 2)

def build_matrix(d):
PS = PS_of_d(d)
CL = CL_of_d(d)
a11 = -(CL + PS) / Vp
a12 = PS / Vp
a21 = PS / Vt
a22 = -(PS / Vt + kdeg)
return a11, a12, a21, a22

# ---------------------------------------------------------------
# 3. Closed-form 2x2 linear system solver (vectorized, no loops)
# ---------------------------------------------------------------
def eig_2x2(a11, a12, a21, a22):
tr = a11 + a22
det = a11 * a22 - a12 * a21
disc = np.clip(tr ** 2 - 4 * det, 0, None) # guard tiny FP noise
sq = np.sqrt(disc)
return (tr + sq) / 2, (tr - sq) / 2

def expm_At(a11, a12, a21, a22, lam1, lam2, tt):
"""Closed-form exp(A t) via Sylvester's formula, broadcast over
particle size (axis 0) and time (axis 1)."""
a11, a12 = a11[:, None], a12[:, None]
a21, a22 = a21[:, None], a22[:, None]
lam1, lam2 = lam1[:, None], lam2[:, None]
tt = tt[None, :]

denom = lam1 - lam2
denom = np.where(np.abs(denom) < 1e-10, 1e-10, denom)
e1, e2 = np.exp(lam1 * tt), np.exp(lam2 * tt)

E11 = ((a11 - lam2) * e1 - (a11 - lam1) * e2) / denom
E12 = (a12 * e1 - a12 * e2) / denom
E21 = (a21 * e1 - a21 * e2) / denom
E22 = ((a22 - lam2) * e1 - (a22 - lam1) * e2) / denom
return E11, E12, E21, E22

def inv_2x2(a11, a12, a21, a22):
det = a11 * a22 - a12 * a21
return a22 / det, -a12 / det, -a21 / det, a11 / det

def simulate_unit_response(d_array):
"""Simulate Cp(t), Ct(t) for a unit infusion rate (1 mg/h),
for every particle size in d_array simultaneously."""
a11, a12, a21, a22 = build_matrix(d_array)
lam1, lam2 = eig_2x2(a11, a12, a21, a22)
inv11, inv12, inv21, inv22 = inv_2x2(a11, a12, a21, a22)
b1 = 1.0 / Vp

Nd = len(d_array)
Cp = np.zeros((Nd, len(t)))
Ct = np.zeros((Nd, len(t)))

# Phase 1: during infusion, x(0)=0, constant input b
# x(t) = A^{-1} (exp(At) - I) b
E11, E12, E21, E22 = expm_At(a11, a12, a21, a22, lam1, lam2, t1)
Mb1, Mb2 = (E11 - 1.0) * b1, E21 * b1
Cp1 = inv11[:, None] * Mb1 + inv12[:, None] * Mb2
Ct1 = inv21[:, None] * Mb1 + inv22[:, None] * Mb2
Cp[:, t1_mask], Ct[:, t1_mask] = Cp1, Ct1

# Phase 2: after infusion stops, free decay from x(Tinf)
x0p, x0t = Cp1[:, -1], Ct1[:, -1]
E11b, E12b, E21b, E22b = expm_At(a11, a12, a21, a22, lam1, lam2, t2)
Cp2 = E11b * x0p[:, None] + E12b * x0t[:, None]
Ct2 = E21b * x0p[:, None] + E22b * x0t[:, None]
Cp[:, ~t1_mask], Ct[:, ~t1_mask] = Cp2, Ct2
return Cp, Ct

def trapz_axis1(y, x):
"""Manual trapezoidal integration along axis 1 (time)."""
dx = np.diff(x)
return np.sum((y[:, :-1] + y[:, 1:]) * dx / 2, axis=1)

# ---------------------------------------------------------------
# 4. Run the sweep over particle size (fully vectorized)
# ---------------------------------------------------------------
d_grid = np.linspace(10, 300, 300)
Cp_unit, Ct_unit = simulate_unit_response(d_grid)

Cp_max_unit = Cp_unit.max(axis=1) # mg/L per unit (mg/h) of infusion
AUC_unit = trapz_axis1(Ct_unit, t) # mg*h/L per unit (mg/h) of infusion

# Because the system is linear in the infusion rate R_in, the largest
# feasible R_in for a given particle size is simply set by the toxicity
# ceiling: R_in*(d) = C_tox / Cp_max_unit(d)
R0_star_d = C_tox / Cp_max_unit
AUC_star_d = R0_star_d * AUC_unit

best_idx = np.nanargmax(AUC_star_d)
d_star = d_grid[best_idx]
R0_star = R0_star_d[best_idx]
AUC_star = AUC_star_d[best_idx]

print(f"Optimal particle diameter : {d_star:.1f} nm")
print(f"Optimal infusion rate : {R0_star:.2f} mg/h")
print(f"Resulting tumor AUC : {AUC_star:.2f} mg*h/L")

# ---------------------------------------------------------------
# 5. Figure 1 - the two competing size-dependent effects
# ---------------------------------------------------------------
fig1, ax1 = plt.subplots(figsize=(8, 5))
ax1.plot(d_grid, PS_of_d(d_grid), color='tab:red', lw=2, label='Permeability PS(d)')
ax1.set_xlabel('Particle diameter (nm)')
ax1.set_ylabel('PS (L/h)', color='tab:red')
ax1.tick_params(axis='y', labelcolor='tab:red')

ax1b = ax1.twinx()
ax1b.plot(d_grid, CL_of_d(d_grid), color='tab:blue', lw=2, label='Clearance CL(d)')
ax1b.set_ylabel('CL (L/h)', color='tab:blue')
ax1b.tick_params(axis='y', labelcolor='tab:blue')

ax1.axvline(d_star, color='k', ls='--', lw=1)
ax1.set_title('EPR permeability vs. plasma clearance as a function of particle size')
fig1.tight_layout()
plt.show()

# ---------------------------------------------------------------
# 6. Figure 2 - time course at the optimal conditions
# ---------------------------------------------------------------
Cp_opt = R0_star * Cp_unit[best_idx]
Ct_opt = R0_star * Ct_unit[best_idx]

fig2, ax2 = plt.subplots(figsize=(8, 5))
ax2.plot(t, Cp_opt, label='Plasma concentration $C_p(t)$', lw=2)
ax2.plot(t, Ct_opt, label='Tumor concentration $C_t(t)$', lw=2)
ax2.axhline(C_tox, color='red', ls=':', label='Toxicity threshold')
ax2.axvline(Tinf, color='gray', ls='--', lw=1, label='End of infusion')
ax2.set_xlabel('Time (h)')
ax2.set_ylabel('Concentration (mg/L)')
ax2.set_title(f'Concentration-time profile at optimum (d={d_star:.0f} nm, R$_{{in}}$={R0_star:.1f} mg/h)')
ax2.legend()
fig2.tight_layout()
plt.show()

# ---------------------------------------------------------------
# 7. Figure 3 - 3D surface: tumor AUC vs particle size & infusion rate
# ---------------------------------------------------------------
R0_grid = np.linspace(0.5, 40, 200)
AUC_2D = np.outer(AUC_unit, R0_grid)
Cpmax_2D = np.outer(Cp_max_unit, R0_grid)
feasible = Cpmax_2D <= C_tox
AUC_feasible = np.where(feasible, AUC_2D, np.nan)

D, R0m = np.meshgrid(d_grid, R0_grid, indexing='ij')

fig3 = plt.figure(figsize=(9, 7))
ax3 = fig3.add_subplot(111, projection='3d')
surf = ax3.plot_surface(D, R0m, AUC_feasible, cmap='viridis',
edgecolor='none', alpha=0.9)
ax3.scatter([d_star], [R0_star], [AUC_star], color='red', s=60, label='Optimum')
ax3.set_xlabel('Particle diameter (nm)')
ax3.set_ylabel('Infusion rate (mg/h)')
ax3.set_zlabel('Tumor AUC (mg*h/L)')
ax3.set_title('Feasible tumor exposure surface (toxicity-constrained region only)')
fig3.colorbar(surf, shrink=0.6, label='Tumor AUC (mg*h/L)')
plt.show()

# ---------------------------------------------------------------
# 8. Figure 4 - contour view with the feasibility boundary
# ---------------------------------------------------------------
fig4, ax4 = plt.subplots(figsize=(8, 6))
cs = ax4.contourf(D, R0m, AUC_feasible, levels=20, cmap='viridis')
ax4.contour(D, R0m, Cpmax_2D, levels=[C_tox], colors='red', linewidths=2)
ax4.plot(d_star, R0_star, 'r*', markersize=18, label='Optimum')
ax4.set_xlabel('Particle diameter (nm)')
ax4.set_ylabel('Infusion rate (mg/h)')
ax4.set_title('Top view: red line = toxicity boundary, star = optimum')
fig4.colorbar(cs, label='Tumor AUC (mg*h/L)')
ax4.legend()
fig4.tight_layout()
plt.show()

# ---------------------------------------------------------------
# 9. Figure 5 - therapeutic efficiency ratio vs particle size
# ---------------------------------------------------------------
fig5, ax5 = plt.subplots(figsize=(8, 5))
ax5.plot(d_grid, AUC_star_d, lw=2, color='tab:green')
ax5.axvline(d_star, color='k', ls='--', lw=1)
ax5.plot(d_star, AUC_star, 'r*', markersize=16)
ax5.set_xlabel('Particle diameter (nm)')
ax5.set_ylabel('Max tumor AUC under toxicity constraint (mg*h/L)')
ax5.set_title('Best achievable tumor exposure at each particle size')
fig5.tight_layout()
plt.show()

Code walkthrough

Parameters (section 1): All values are illustrative, chosen to produce realistic-looking pharmacokinetic behavior rather than sourced from a specific real drug. Vt is deliberately tiny (10 mL) compared to Vp (5 L), which is what makes drug exchange with the tumor compartment fast relative to systemic elimination — a common feature of small, well-perfused tumor lesions.

Size-dependent physiology (section 2): PS_of_d encodes the EPR effect as a Gaussian peaking at 100 nm — permeability collapses quickly for particles that are too small or too large. CL_of_d encodes the trade-off on the clearance side: a parabola with a minimum at 130 nm, reflecting that very small particles are lost through renal filtration and very large ones are captured by the liver/spleen. These two peaks don’t coincide, which is exactly what creates a genuine, non-trivial optimum — not just an assumption stated up front.

Closed-form linear solver (section 3): Instead of stepping the ODE forward in time with an integrator, eig_2x2 and expm_At compute the matrix exponential $e^{At}$ analytically via Sylvester’s formula for a 2×2 matrix. All of the array shapes carry a leading “particle size” axis, so expm_At returns the exponential for all 300 particle sizes and every time point in one shot — this is the core trick that keeps everything fast. simulate_unit_response then assembles the full concentration trajectories for both the infusion phase and the post-infusion decay phase, again for every particle size simultaneously.

The sweep (section 4): We only ever solve the model once, at a unit infusion rate. Because the whole system is linear, the response to any infusion rate is just that unit response scaled by R0. This lets us derive the toxicity-constrained optimal infusion rate for every particle size analytically: R0_star_d = C_tox / Cp_max_unit, with no additional simulation required. Maximizing AUC_star_d over the particle-size axis then gives the overall optimum.

Figures 3 and 4 still sweep a full 200-point grid of infusion rates for the visualization — this is just an outer product (np.outer), essentially free computationally, and it lets the reader see the entire feasible operating region, not just the single optimal point.

Reading the results

Running the script gives:

1
2
3
Optimal particle diameter : 127.4 nm
Optimal infusion rate : 27.71 mg/h
Resulting tumor AUC : 104.91 mg*h/L

The optimal size sits between the pure permeability peak (100 nm) and the pure circulation-time peak (130 nm) — the model is genuinely balancing two competing effects rather than just reproducing one of the input assumptions. About 82% of the sampled (size, dose) grid turns out to be within the toxicity constraint, but only a narrow ridge along that constraint boundary actually maximizes tumor exposure — which is exactly what Figures 3 and 4 are meant to make visible.

Figure 1 shows why 127 nm is the answer: permeability is already past its peak and declining, but clearance is still improving (still falling), and the net effect nets out slightly past the permeability optimum.

Figure 2 shows the actual concentration-time trace at the optimum — plasma concentration rises during the 2-hour infusion and decays afterward, while tumor concentration lags slightly behind and clears more slowly, since it’s driven by the smaller kdeg rate rather than plasma clearance.

Figure 3 is the main result: a 3D surface of tumor AUC over the full (particle size, infusion rate) plane, with everything above the toxicity ceiling masked out (NaN) so only physically realizable treatment plans are shown. The red marker sits on the ridge of the feasible region where AUC is maximized.

Figure 4 is the same surface from directly above, as a contour plot, with the toxicity boundary drawn explicitly in red — this is the easiest view for reading off exact coordinates of the optimum.

Figure 5 collapses the whole problem down to a single curve: for each particle size, what is the best tumor exposure achievable without crossing the toxicity limit? This makes the interior optimum at ~127 nm unambiguous.

Takeaway

Framing drug delivery as a small linear compartmental model, rather than reaching for a general-purpose ODE solver, buys two things at once: a rigorous, closed-form answer to “what particle size and dose are optimal under a safety constraint,” and a computation cheap enough to sweep and visualize the entire design space rather than just reporting a single number.

Optimizing Genome Analysis Models

Predicting Treatment Efficacy from Genetic Mutations

Precision medicine hinges on one central question: given a patient’s genetic mutation profile, how well will a specific treatment work? Building a model that can answer this reliably requires more than just throwing data at an algorithm — it requires careful feature engineering, systematic hyperparameter optimization, and rigorous validation. In this article, we’ll build a complete pipeline that predicts treatment efficacy from simulated mutation data, optimize it using a grid search across two key hyperparameters, and visualize the entire optimization landscape in 3D.

The Problem Setup

We model treatment efficacy as a continuous score $y \in [0, 100]$ that depends on a patient’s mutation profile $\mathbf{x} = (x_1, x_2, \dots, x_n)$, where each $x_i \in {0, 1}$ indicates the presence or absence of a mutation in gene $i$. The true underlying relationship is rarely linear — genes interact with each other (epistasis), and some mutations only matter in combination with others. We express this as:

$$
y = f(\mathbf{x}) + \epsilon, \quad f(\mathbf{x}) = \beta_0 + \sum_{i=1}^{n} \beta_i x_i + \sum_{i<j} \gamma_{ij} x_i x_j
$$

where $\beta_i$ are main effects, $\gamma_{ij}$ are interaction effects, and $\epsilon \sim \mathcal{N}(0, \sigma^2)$ is measurement noise. Our goal is to recover an approximation of $f$ using a machine learning model trained only on observed $(\mathbf{x}, y)$ pairs.

The Full Source Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
# =========================================================
# Genome Analysis Model Optimization
# Predicting Treatment Efficacy from Genetic Mutation Profiles
# =========================================================

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import r2_score, mean_squared_error

# ---------------------------------------------------------
# 1. Reproducibility
# ---------------------------------------------------------
RANDOM_STATE = 42
np.random.seed(RANDOM_STATE)

# ---------------------------------------------------------
# 2. Synthetic genomic mutation dataset generation
# ---------------------------------------------------------
N_PATIENTS = 1200
N_GENES = 20
GENE_NAMES = [f"GENE_{i+1:02d}" for i in range(N_GENES)]

# Binary mutation matrix: 1 = mutation present, 0 = absent
mutation_prob = np.random.uniform(0.1, 0.4, size=N_GENES)
X = np.random.binomial(1, mutation_prob, size=(N_PATIENTS, N_GENES)).astype(float)

# True (hidden) main effects for each gene
true_main_effects = np.random.normal(loc=0, scale=6, size=N_GENES)

# A handful of genuine gene-gene interactions (epistasis)
interaction_pairs = [(0, 1), (2, 5), (3, 7), (8, 9), (10, 15)]
interaction_strengths = [8.0, -6.5, 5.0, -4.0, 7.5]

baseline_efficacy = 45.0
y = np.full(N_PATIENTS, baseline_efficacy)

for i in range(N_GENES):
y += true_main_effects[i] * X[:, i]

for (gi, gj), strength in zip(interaction_pairs, interaction_strengths):
y += strength * X[:, gi] * X[:, gj]

# Measurement noise
noise = np.random.normal(0, 5.0, size=N_PATIENTS)
y += noise
y = np.clip(y, 0, 100) # efficacy score bounded between 0 and 100

df = pd.DataFrame(X, columns=GENE_NAMES)
df["treatment_efficacy"] = y

print("Dataset shape:", df.shape)
print(df.head())

# ---------------------------------------------------------
# 3. Train / test split
# ---------------------------------------------------------
X_train, X_test, y_train, y_test = train_test_split(
df[GENE_NAMES], df["treatment_efficacy"],
test_size=0.25, random_state=RANDOM_STATE
)

# ---------------------------------------------------------
# 4. Hyperparameter optimization landscape (n_estimators x max_depth)
# ---------------------------------------------------------
n_estimators_range = [20, 60, 100, 140, 180]
max_depth_range = [2, 4, 6, 8, 10]

score_grid = np.zeros((len(max_depth_range), len(n_estimators_range)))

for di, depth in enumerate(max_depth_range):
for ni, n_est in enumerate(n_estimators_range):
model = RandomForestRegressor(
n_estimators=n_est,
max_depth=depth,
random_state=RANDOM_STATE,
n_jobs=-1
)
model.fit(X_train, y_train)
preds = model.predict(X_test)
score_grid[di, ni] = r2_score(y_test, preds)

# ---------------------------------------------------------
# 5. Fine-grained optimization with GridSearchCV
# ---------------------------------------------------------
param_grid = {
"n_estimators": [80, 100, 120, 150, 180, 200],
"max_depth": [4, 6, 8, 10, 12, None],
"min_samples_leaf": [1, 2, 4]
}

grid_search = GridSearchCV(
estimator=RandomForestRegressor(random_state=RANDOM_STATE, n_jobs=-1),
param_grid=param_grid,
cv=5,
scoring="r2",
n_jobs=-1
)
grid_search.fit(X_train, y_train)

best_model = grid_search.best_estimator_
print("\nBest hyperparameters found:", grid_search.best_params_)

# ---------------------------------------------------------
# 6. Final evaluation
# ---------------------------------------------------------
y_pred = best_model.predict(X_test)
final_r2 = r2_score(y_test, y_pred)
final_mse = mean_squared_error(y_test, y_pred)

print(f"\nFinal R^2 score : {final_r2:.4f}")
print(f"Final MSE : {final_mse:.4f}")

# ---------------------------------------------------------
# 7. Visualization 1: Actual vs Predicted Efficacy
# ---------------------------------------------------------
plt.figure(figsize=(7, 7))
plt.scatter(y_test, y_pred, alpha=0.5, edgecolor="k", s=40, color="#2E86AB")
lims = [0, 100]
plt.plot(lims, lims, "r--", linewidth=2, label="Perfect prediction")
plt.xlabel("Actual Treatment Efficacy")
plt.ylabel("Predicted Treatment Efficacy")
plt.title(f"Actual vs Predicted Efficacy (R² = {final_r2:.3f})")
plt.legend()
plt.grid(alpha=0.3)
plt.tight_layout()
plt.show()

# ---------------------------------------------------------
# 8. Visualization 2: Feature Importance
# ---------------------------------------------------------
importances = best_model.feature_importances_
sorted_idx = np.argsort(importances)[::-1]

plt.figure(figsize=(9, 8))
plt.barh(
[GENE_NAMES[i] for i in sorted_idx][::-1],
importances[sorted_idx][::-1],
color="#A23B72"
)
plt.xlabel("Feature Importance")
plt.title("Gene Contribution to Treatment Efficacy Prediction")
plt.tight_layout()
plt.show()

# ---------------------------------------------------------
# 9. Visualization 3: 3D Hyperparameter Optimization Surface
# ---------------------------------------------------------
X_grid, Y_grid = np.meshgrid(n_estimators_range, max_depth_range)

fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection="3d")
surf = ax.plot_surface(
X_grid, Y_grid, score_grid,
cmap="viridis", edgecolor="k", linewidth=0.3, alpha=0.9
)
ax.set_xlabel("n_estimators")
ax.set_ylabel("max_depth")
ax.set_zlabel("R² Score")
ax.set_title("Hyperparameter Optimization Landscape")
fig.colorbar(surf, shrink=0.6, aspect=12, label="R² Score")
plt.tight_layout()
plt.show()
Dataset shape: (1200, 21)
   GENE_01  GENE_02  GENE_03  GENE_04  GENE_05  GENE_06  GENE_07  GENE_08  \
0      0.0      0.0      0.0      0.0      0.0      0.0      0.0      0.0   
1      0.0      0.0      0.0      1.0      0.0      0.0      0.0      0.0   
2      0.0      0.0      1.0      0.0      0.0      0.0      0.0      1.0   
3      1.0      1.0      0.0      0.0      0.0      0.0      0.0      0.0   
4      0.0      1.0      0.0      0.0      1.0      0.0      0.0      1.0   

   GENE_09  GENE_10  ...  GENE_12  GENE_13  GENE_14  GENE_15  GENE_16  \
0      0.0      0.0  ...      0.0      0.0      1.0      1.0      0.0   
1      0.0      0.0  ...      1.0      1.0      1.0      0.0      1.0   
2      0.0      1.0  ...      0.0      0.0      0.0      0.0      0.0   
3      1.0      0.0  ...      1.0      1.0      0.0      0.0      0.0   
4      0.0      0.0  ...      0.0      1.0      0.0      0.0      1.0   

   GENE_17  GENE_18  GENE_19  GENE_20  treatment_efficacy  
0      0.0      0.0      0.0      0.0           30.889048  
1      0.0      0.0      0.0      0.0           42.704561  
2      0.0      0.0      0.0      0.0           35.221280  
3      0.0      0.0      0.0      0.0           20.029387  
4      0.0      0.0      1.0      0.0           20.054577  

[5 rows x 21 columns]

Best hyperparameters found: {'max_depth': 12, 'min_samples_leaf': 2, 'n_estimators': 200}

Final R^2 score : 0.7253
Final MSE        : 46.9035

Code Walkthrough

Step 1–2: Building a realistic synthetic dataset. Since real clinical genomic datasets are sensitive and hard to share, we simulate one that mimics real biological structure. Each of the 20 genes has an independent mutation probability drawn from np.random.uniform(0.1, 0.4), so mutation frequency varies across genes — just as it does in real cohorts. We then assign each gene a hidden “true” effect size (true_main_effects) drawn from a normal distribution, meaning some mutations boost efficacy and others suppress it.

Critically, we also inject five interaction pairs — combinations of two genes whose joint presence has an effect beyond their individual contributions. This is what makes the problem non-trivial: a purely linear model would systematically miss these epistatic effects, while a tree-based ensemble like Random Forest can capture them naturally through recursive splitting.

Step 3: Train/test split. We hold out 25% of patients purely for unbiased evaluation, ensuring the model is judged on data it has never seen.

Step 4: Building the 3D optimization surface. Before running a formal grid search, we manually sweep two of the most influential Random Forest hyperparameters — n_estimators (number of trees) and max_depth (tree depth) — training a fresh model at every combination and recording its $R^2$ score on the test set. This produces a 5×5 grid of scores that we can later render as a 3D surface, giving an intuitive picture of how performance responds to these two knobs simultaneously.

Step 5: Fine-grained optimization with GridSearchCV. This is the real optimization step. We expand the search to three hyperparameters (n_estimators, max_depth, min_samples_leaf) and use 5-fold cross-validation so each candidate configuration is evaluated on five different train/validation splits rather than just one, reducing the risk of overfitting to a lucky split. GridSearchCV automatically returns the best-performing configuration via best_params_.

Step 6: Final evaluation. We compute two standard regression metrics on the held-out test set:

$R^2$ tells us the proportion of variance in treatment efficacy explained by the model (closer to 1 is better), while MSE quantifies the average squared prediction error in the original efficacy units.

Speeding Things Up

The 5×5 manual sweep in Step 4 trains 25 separate Random Forests sequentially, and GridSearchCV in Step 5 evaluates $6 \times 6 \times 3 \times 5 = 540$ model fits. Two things already keep this fast in the code above: n_jobs=-1 parallelizes both the individual Random Forest’s tree-building and GridSearchCV‘s cross-validation folds across all available CPU cores. If you’re working with a larger gene panel (hundreds of genes) or more patients, you can accelerate further by:

  • Replacing the exhaustive GridSearchCV with RandomizedSearchCV (samples a fixed number of random combinations instead of testing every single one), which cuts search time dramatically with minimal loss in solution quality.
  • Reducing cv from 5 to 3 folds during the initial coarse search, then refining with 5-fold cross-validation only around the best region.
  • Using HistGradientBoostingRegressor from sklearn.ensemble, which is substantially faster than RandomForestRegressor on larger tabular datasets while achieving comparable accuracy.

Interpreting the Visualizations

Actual vs Predicted scatter plot. Points clustered tightly around the red diagonal line indicate accurate predictions; the further a point strays from that line, the larger the prediction error for that patient. The $R^2$ value in the title summarizes this visually as a single number.

Feature importance chart. This bar chart reveals which genes the model relies on most heavily when making predictions. Genes involved in our simulated interaction pairs (e.g., GENE_01, GENE_02, GENE_09, GENE_10) should rank prominently, since their combined effect creates a strong, learnable signal — even though Random Forest never sees the interaction terms explicitly, it discovers them through splits on both features.

3D hyperparameter optimization surface. This is the most revealing plot for understanding the optimization process itself. The x-axis shows n_estimators, the y-axis shows max_depth, and the z-axis (height and color) shows the resulting $R^2$ score. Typically you’ll observe the surface rise sharply as max_depth increases from very shallow trees (which underfit and can’t capture interactions), then plateau once trees are deep enough to model the epistatic relationships, while increasing n_estimators provides diminishing but steady improvements by reducing variance through averaging. Flat regions at the top of the surface indicate that further increasing either parameter yields little additional benefit — useful information for choosing an efficient, not just an accurate, final model.

Closing Thoughts

This pipeline demonstrates the full lifecycle of a genome-to-efficacy prediction model: simulating biologically plausible mutation data with hidden gene-gene interactions, training an ensemble model capable of capturing non-linear epistatic effects, systematically optimizing hyperparameters through both an exploratory 3D landscape sweep and a rigorous cross-validated grid search, and finally validating and visualizing the results. The same structure scales naturally to real-world genomic datasets — swap in actual variant call data, expand the gene panel, and the optimization machinery here requires no changes.

Parameter Optimization for Diagnostic Imaging AI

Sharpening Cancer Detection Accuracy

Why threshold and weight calibration matters more than model architecture

A modern diagnostic imaging AI rarely fails because its underlying network is too weak. It fails because the operating point — the combination of feature weights and decision threshold that turns a continuous malignancy score into a binary “biopsy / no biopsy” recommendation — is miscalibrated. In cancer screening, this operating point decision is a constrained optimization problem: clinicians want to push sensitivity (catch every malignant case) as high as possible, but not at the cost of collapsing specificity and flooding the pipeline with false positives.

This is a textbook Neyman–Pearson trade-off, and it can be formulated as a differentiable optimization problem over the model’s weight vector, bias, and decision threshold. Below, we build a radiomics-style scoring model, derive a smooth surrogate objective for the (non-differentiable) sensitivity/specificity trade-off, and optimize it end-to-end with a hand-rolled, fully vectorized Adam optimizer.

Mathematical formulation

Each lesion is represented by a radiomic feature vector $x_i \in \mathbb{R}^d$ (texture entropy, margin spiculation, intensity variance, and so on). A linear score is passed through a sigmoid to produce a malignancy probability:

$$
z_i = w^\top x_i + b, \qquad s_i = \sigma(z_i) = \frac{1}{1+e^{-z_i}}
$$

A sample is flagged malignant when $s_i > \tau$. Sensitivity and specificity are step functions of $\tau$, which have zero gradient almost everywhere. We replace the hard indicator with a steep sigmoid of steepness $k$:

$$
\widehat{\text{Sens}}(w,b,\tau) = \frac{1}{|\mathcal{M}|}\sum_{i \in \mathcal{M}} \sigma\big(k(s_i - \tau)\big), \qquad
\widehat{\text{Spec}}(w,b,\tau) = \frac{1}{|\mathcal{N}|}\sum_{i \in \mathcal{N}} \sigma\big(k(\tau - s_i)\big)
$$

where $\mathcal{M}$ and $\mathcal{N}$ are the malignant and benign index sets. The clinical objective — weighted toward recall, since missed cancers are costlier than false alarms — is:

$$
L(w,b,\tau) = \alpha ,\widehat{\text{Sens}} + (1-\alpha),\widehat{\text{Spec}} - \lambda |w|_2^2
$$

with $\alpha \in (0.5, 1)$ biasing the trade-off toward sensitivity, and $\lambda$ an $L_2$ regularizer preventing weight blow-up. The gradients follow from the chain rule through $s_i = \sigma(z_i)$:

$$
\frac{\partial L}{\partial w} = \sum_i \frac{\partial L}{\partial s_i}, s_i(1-s_i), x_i ;-; 2\lambda w, \qquad
\frac{\partial L}{\partial b} = \sum_i \frac{\partial L}{\partial s_i}, s_i(1-s_i)
$$

$$
\frac{\partial L}{\partial \tau} = -\alpha k \cdot \overline{g_{\mathcal{M}}} + (1-\alpha)k \cdot \overline{g_{\mathcal{N}}}
$$

where $g_i = \sigma\big(k(\cdot)\big)\big(1-\sigma(k(\cdot))\big)$ is the local sigmoid derivative. Every term above is expressed as a sum over the dataset, which means the whole gradient computation vectorizes into pure NumPy matrix operations — no per-sample Python loops, and no risk of slow execution even for a dense hyperparameter landscape scan.

Full source code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from sklearn.metrics import roc_curve, auc
from sklearn.model_selection import train_test_split

plt.style.use('dark_background')
np.random.seed(42)

# ------------------------------------------------------------------
# 1. Synthetic radiomics dataset generation
# ------------------------------------------------------------------
feature_names = ['Texture Entropy', 'Margin Spiculation', 'Intensity Variance',
'Shape Irregularity', 'Vascularity Index', 'Microcalc. Score']
d = len(feature_names)

n_benign = 400
n_malignant = 300

mean_benign = np.zeros(d)
mean_malignant = np.array([1.4, 1.1, 0.9, 1.3, 0.6, 1.0])
cov_base = 0.6 * np.eye(d) + 0.15 * (np.ones((d, d)) - np.eye(d))

X_benign = np.random.multivariate_normal(mean_benign, cov_base, size=n_benign)
X_malignant = np.random.multivariate_normal(mean_malignant, cov_base, size=n_malignant)

X_raw = np.vstack([X_benign, X_malignant])
y = np.concatenate([np.zeros(n_benign), np.ones(n_malignant)])

X_train_raw, X_test_raw, y_train, y_test = train_test_split(
X_raw, y, test_size=0.3, random_state=42, stratify=y)

mu = X_train_raw.mean(axis=0)
sigma_ = X_train_raw.std(axis=0)
X_train = (X_train_raw - mu) / sigma_
X_test = (X_test_raw - mu) / sigma_

# ------------------------------------------------------------------
# 2. Differentiable sensitivity / specificity surrogate objective
# ------------------------------------------------------------------
def sigmoid(z):
z = np.clip(z, -30, 30)
return 1.0 / (1.0 + np.exp(-z))

def objective_and_grad(w, b, tau, X, y, k=12.0, alpha=0.6, lam=0.01):
z = X @ w + b
s = sigmoid(z)
ds = s * (1.0 - s)

mask_pos = (y == 1)
mask_neg = (y == 0)
m = mask_pos.sum()
n = mask_neg.sum()

sig_pos = sigmoid(k * (s[mask_pos] - tau))
sig_neg = sigmoid(k * (tau - s[mask_neg]))

soft_sens = sig_pos.mean()
soft_spec = sig_neg.mean()
loss = alpha * soft_sens + (1 - alpha) * soft_spec - lam * np.sum(w ** 2)

g_pos = sig_pos * (1 - sig_pos)
g_neg = sig_neg * (1 - sig_neg)

dL_ds = np.zeros_like(s)
dL_ds[mask_pos] = alpha * k * g_pos / m
dL_ds[mask_neg] += (1 - alpha) * (-k) * g_neg / n

dL_dz = dL_ds * ds
grad_w = X.T @ dL_dz - 2 * lam * w
grad_b = dL_dz.sum()
grad_tau = -alpha * k * (g_pos / m).sum() + (1 - alpha) * k * (g_neg / n).sum()

return loss, grad_w, grad_b, grad_tau, soft_sens, soft_spec

# ------------------------------------------------------------------
# 3. Adam optimizer (fully vectorized, no per-sample loops)
# ------------------------------------------------------------------
def train(X, y, iters=600, lr=0.08, k=12.0, alpha=0.6, lam=0.01):
d = X.shape[1]
w = np.random.randn(d) * 0.1
b = 0.0
tau = 0.5

mw = np.zeros(d); vw = np.zeros(d)
mb = vb = mt = vt = 0.0
beta1, beta2, eps = 0.9, 0.999, 1e-8

history = {'loss': [], 'sens': [], 'spec': []}

for t in range(1, iters + 1):
loss, gw, gb, gt, sens, spec = objective_and_grad(w, b, tau, X, y, k, alpha, lam)

mw = beta1 * mw + (1 - beta1) * gw
vw = beta2 * vw + (1 - beta2) * (gw ** 2)
w += lr * (mw / (1 - beta1 ** t)) / (np.sqrt(vw / (1 - beta2 ** t)) + eps)

mb = beta1 * mb + (1 - beta1) * gb
vb = beta2 * vb + (1 - beta2) * (gb ** 2)
b += lr * (mb / (1 - beta1 ** t)) / (np.sqrt(vb / (1 - beta2 ** t)) + eps)

mt = beta1 * mt + (1 - beta1) * gt
vt = beta2 * vt + (1 - beta2) * (gt ** 2)
tau += lr * (mt / (1 - beta1 ** t)) / (np.sqrt(vt / (1 - beta2 ** t)) + eps)
tau = np.clip(tau, 0.05, 0.95)

history['loss'].append(loss)
history['sens'].append(sens)
history['spec'].append(spec)

return w, b, tau, history

w_opt, b_opt, tau_opt, history = train(X_train, y_train)

rng_baseline = np.random.RandomState(0)
w_base = rng_baseline.randn(d) * 0.1
b_base = 0.0

# ------------------------------------------------------------------
# 4. Evaluation on held-out test set
# ------------------------------------------------------------------
def hard_metrics(w, b, tau, X, y):
s = sigmoid(X @ w + b)
pred = (s > tau).astype(int)
tp = np.sum((pred == 1) & (y == 1))
fn = np.sum((pred == 0) & (y == 1))
tn = np.sum((pred == 0) & (y == 0))
fp = np.sum((pred == 1) & (y == 0))
sens = tp / (tp + fn + 1e-12)
spec = tn / (tn + fp + 1e-12)
return sens, spec, s

sens_opt, spec_opt, s_test_opt = hard_metrics(w_opt, b_opt, tau_opt, X_test, y_test)
sens_base, spec_base, s_test_base = hard_metrics(w_base, b_base, 0.5, X_test, y_test)

fpr_opt, tpr_opt, _ = roc_curve(y_test, s_test_opt)
auc_opt = auc(fpr_opt, tpr_opt)
fpr_base, tpr_base, _ = roc_curve(y_test, s_test_base)
auc_base = auc(fpr_base, tpr_base)

print(f"Optimized -> Sensitivity: {sens_opt:.3f}, Specificity: {spec_opt:.3f}, AUC: {auc_opt:.3f}")
print(f"Baseline -> Sensitivity: {sens_base:.3f}, Specificity: {spec_base:.3f}, AUC: {auc_base:.3f}")
print(f"Optimized tau: {tau_opt:.3f}")
print(f"Optimized w: {np.round(w_opt, 3)}")

# ------------------------------------------------------------------
# 5. Visualization
# ------------------------------------------------------------------
fig1, axes = plt.subplots(1, 2, figsize=(13, 5))
axes[0].plot(history['loss'], color='#00e5ff', lw=2)
axes[0].set_title('Objective value during optimization', color='white')
axes[0].set_xlabel('Iteration'); axes[0].set_ylabel('L(w, b, tau)')
axes[0].grid(alpha=0.2)

axes[1].plot(history['sens'], color='#ff5c8a', lw=2, label='Soft sensitivity')
axes[1].plot(history['spec'], color='#7dff8a', lw=2, label='Soft specificity')
axes[1].set_title('Sensitivity / Specificity convergence', color='white')
axes[1].set_xlabel('Iteration'); axes[1].legend()
axes[1].grid(alpha=0.2)
plt.tight_layout()
plt.show()

fig2 = plt.figure(figsize=(6.5, 6))
ax2 = fig2.add_subplot(111)
ax2.plot(fpr_base, tpr_base, color='#888888', lw=2, ls='--',
label=f'Baseline (AUC = {auc_base:.3f})')
ax2.plot(fpr_opt, tpr_opt, color='#00e5ff', lw=2.5,
label=f'Optimized (AUC = {auc_opt:.3f})')
ax2.plot([0, 1], [0, 1], color='gray', lw=1, ls=':')
ax2.set_xlabel('False Positive Rate'); ax2.set_ylabel('True Positive Rate')
ax2.set_title('ROC curve: before vs after optimization', color='white')
ax2.legend(loc='lower right')
ax2.grid(alpha=0.2)
plt.tight_layout()
plt.show()

direction = w_opt / np.linalg.norm(w_opt)
proj = X_train @ direction

scale_grid = np.linspace(0.0, 2.0 * np.linalg.norm(w_opt), 50)
tau_grid = np.linspace(0.05, 0.95, 50)
SCALE, TAU = np.meshgrid(scale_grid, tau_grid)

proj_pos = proj[y_train == 1]
proj_neg = proj[y_train == 0]
k_land, alpha_land, lam_land = 12.0, 0.6, 0.01

Z_pos = SCALE[:, :, None] * proj_pos[None, None, :] + b_opt
soft_sens_grid = sigmoid(k_land * (sigmoid(Z_pos) - TAU[:, :, None])).mean(axis=2)

Z_neg = SCALE[:, :, None] * proj_neg[None, None, :] + b_opt
soft_spec_grid = sigmoid(k_land * (TAU[:, :, None] - sigmoid(Z_neg))).mean(axis=2)

L_grid = (alpha_land * soft_sens_grid + (1 - alpha_land) * soft_spec_grid
- lam_land * (SCALE ** 2))

fig3 = plt.figure(figsize=(8, 6.5))
ax3 = fig3.add_subplot(111, projection='3d')
surf = ax3.plot_surface(SCALE, TAU, L_grid, cmap='plasma',
linewidth=0, antialiased=True, alpha=0.95)
best_idx = np.unravel_index(np.argmax(L_grid), L_grid.shape)
ax3.scatter([SCALE[best_idx]], [TAU[best_idx]], [L_grid[best_idx]],
color='white', s=60, marker='*', label='Grid optimum')
ax3.set_xlabel('||w|| (weight scale)')
ax3.set_ylabel('tau (threshold)')
ax3.set_zlabel('L(w,b,tau)')
ax3.set_title('Optimization landscape', color='white')
fig3.colorbar(surf, shrink=0.6, pad=0.1)
ax3.legend()
plt.tight_layout()
plt.show()

top3_idx = np.argsort(-np.abs(w_opt))[:3]
names3 = [feature_names[i] for i in top3_idx]
w3 = w_opt[top3_idx].copy()
if abs(w3[2]) < 1e-6:
w3[2] = 1e-6

fig4 = plt.figure(figsize=(8, 6.5))
ax4 = fig4.add_subplot(111, projection='3d')
mal = y_train == 1
ben = y_train == 0
ax4.scatter(X_train[ben, top3_idx[0]], X_train[ben, top3_idx[1]], X_train[ben, top3_idx[2]],
color='#7dff8a', s=14, alpha=0.6, label='Benign')
ax4.scatter(X_train[mal, top3_idx[0]], X_train[mal, top3_idx[1]], X_train[mal, top3_idx[2]],
color='#ff5c8a', s=14, alpha=0.6, label='Malignant')

x1_range = np.linspace(X_train[:, top3_idx[0]].min(), X_train[:, top3_idx[0]].max(), 20)
x2_range = np.linspace(X_train[:, top3_idx[1]].min(), X_train[:, top3_idx[1]].max(), 20)
X1G, X2G = np.meshgrid(x1_range, x2_range)
logit_tau = np.log(tau_opt / (1 - tau_opt))
X3G = (logit_tau - b_opt - w3[0] * X1G - w3[1] * X2G) / w3[2]

ax4.plot_surface(X1G, X2G, X3G, color='#00e5ff', alpha=0.25, linewidth=0)
ax4.set_xlabel(names3[0]); ax4.set_ylabel(names3[1]); ax4.set_zlabel(names3[2])
ax4.set_title('Decision boundary in top-3 feature space', color='white')
ax4.legend()
plt.tight_layout()
plt.show()
Optimized -> Sensitivity: 0.922, Specificity: 0.850, AUC: 0.967
Baseline  -> Sensitivity: 0.878, Specificity: 0.800, AUC: 0.924
Optimized tau: 0.476
Optimized w: [0.853 0.473 0.176 0.784 0.186 0.519]

Code walkthrough

Dataset generation. Benign and malignant lesions are drawn from correlated multivariate normal distributions over six radiomic-style features. The malignant class is shifted along several axes, mimicking how real malignant lesions tend to show elevated texture entropy, margin spiculation, and shape irregularity. Standardization statistics are computed strictly from the training split to avoid leaking test-set information — a common pitfall in medical ML pipelines.

Surrogate objective. objective_and_grad computes the loss and all three gradients ($\partial L/\partial w$, $\partial L/\partial b$, $\partial L/\partial \tau$) in one pass using boolean masking and matrix multiplication. There is no Python-level iteration over samples; X.T @ dL_dz performs the entire weighted feature aggregation in a single BLAS call.

Optimizer. A hand-written Adam update is used instead of scipy.optimize, because Adam’s per-parameter adaptive learning rates handle the very different curvature scales of $w$ (six-dimensional, small gradients) and $\tau$ (one-dimensional, sharp gradients near the decision boundary) gracefully without manual tuning.

Evaluation. hard_metrics reverts to the true, non-smoothed sensitivity/specificity definition on the untouched test set, giving an honest performance readout rather than the optimistic soft-surrogate value.

Landscape computation. Rather than looping over a 50×50 grid of $(|w|, \tau)$ combinations in Python, the grid is broadcast against all training projections at once via NumPy’s [:, :, None] trick, producing a $(50,50,n)$ tensor collapsed with .mean(axis=2). This evaluates 122,500 loss values without a single explicit loop, keeping runtime on the order of a second even on Colab’s default CPU runtime.

Reading the graphs

Convergence plot. The left panel shows the composite objective $L(w,b,\tau)$ climbing monotonically as Adam updates the parameters — confirming the gradient derivation is correct and the surrogate is smooth enough for stable ascent. The right panel separates this into soft sensitivity and soft specificity individually, showing the trade-off being resolved: sensitivity typically rises faster early on since $\alpha = 0.6$ weights it more heavily, while specificity stabilizes at a lower but still acceptable plateau.

ROC comparison. Plotting the ROC curve for the untrained baseline weights against the optimized weights isolates the effect of the optimization procedure itself, independent of threshold choice (ROC curves are threshold-agnostic by construction). A visibly larger AUC for the optimized curve demonstrates that the learned feature weighting genuinely separates malignant from benign scores better, not merely that a lucky threshold was picked.

3D optimization landscape. This surface sweeps over the overall weight-vector scale and the decision threshold $\tau$, holding the weight direction fixed at its optimized value. The peak of this surface (marked with a white star) shows where the composite objective is maximized jointly over scale and threshold — visualizing that too small a weight scale under-separates the classes (flattening both sensitivity and specificity), while too aggressive a threshold trades one metric for the other. This is the clearest visual evidence of why a specific $(|w|, \tau)$ pair is optimal rather than arbitrary.

3D decision boundary in feature space. The three radiomic features with the largest optimized weight magnitude are plotted directly, with benign and malignant lesions colored separately. The translucent plane is the exact decision boundary the optimized model uses in this subspace — solved analytically from $w^\top x + b = \text{logit}(\tau)$ while holding the remaining (standardized, zero-mean) features fixed. Seeing malignant points cluster on one side of this plane gives an intuitive, spatial confirmation that the optimized parameters correspond to a geometrically sensible separating boundary, not just an abstract numerical improvement.

Closing notes

The same surrogate-objective-plus-Adam pattern generalizes directly to real deployed models: swap the linear score for a frozen CNN’s final logit layer, keep $w$ and $b$ as a lightweight recalibration head, and the identical gradient derivation still applies. This is, in essence, how many clinical-grade imaging AI systems perform post-hoc threshold calibration against a target sensitivity constraint without retraining the entire network.

Optimizing a Limited Cancer-Screening Budget

A Practical Python Approach

Public health agencies never have unlimited money. Every year, a health ministry or insurer has to decide how many dollars go toward breast cancer screening, how many toward colorectal screening, how many toward lung cancer screening, and so on — knowing that each additional dollar catches slightly fewer new cases than the one before it. This is a classic concave resource allocation problem, and it’s a great case study for combining optimization theory with Python.

In this article we’ll build a small but realistic model, solve it two different ways (a general-purpose optimizer and a much faster specialized algorithm), and visualize the results — including a 3D view of the benefit landscape.

The core idea: diminishing returns

Each screening program has a saturation point. If you screen almost everyone in a population, you eventually catch nearly every detectable case, and additional spending buys you very little extra benefit. This is naturally modeled with an exponential saturation curve:

$$
B_i(x_i) = a_i \left(1 - e^{-x_i / b_i}\right)
$$

Here $x_i$ is the budget (in units of $10,000) given to program $i$, $a_i$ is the maximum number of cases that program could ever detect (its ceiling), and $b_i$ controls how quickly it approaches that ceiling.

The optimization problem is:

$$
\max_{x_1,\dots,x_n} \sum_{i=1}^{n} a_i\left(1 - e^{-x_i/b_i}\right) \quad \text{subject to} \quad \sum_{i=1}^n x_i = X_{\text{total}}, \quad x_i \ge 0
$$

Because every $B_i$ is concave, this problem has a beautiful classical solution: at the optimum, the marginal benefit of every program must be equal. This is the same “water-filling” principle used in economics and information theory. Formally, there exists a single multiplier $\lambda$ such that:

$$
\frac{dB_i}{dx_i} = \frac{a_i}{b_i} e^{-x_i/b_i} = \lambda \quad \text{for every program with } x_i > 0
$$

Solving this equation for $x_i$ gives a closed form:

$$
x_i(\lambda) = \max\left(0,\ -b_i \ln!\left(\frac{\lambda b_i}{a_i}\right)\right)
$$

and we only need to search for the $\lambda$ that makes $\sum_i x_i(\lambda) = X_{\text{total}}$. That single insight is what lets us build a much faster solver than a generic optimizer.

The concrete example

Let’s imagine a regional health authority with four screening programs and a total annual budget of $1,000,000 (100 units of $10,000).

Program Max detectable cases ($a_i$) Diminishing-returns scale ($b_i$)
Breast 120 25
Colorectal 90 15
Lung 150 40
Cervical 60 10

These numbers are illustrative but structured to reflect reality: lung cancer screening has a high ceiling but also needs more spending to get there (low-dose CT is expensive), while cervical screening saturates quickly with modest spending (Pap smears are cheap and coverage is often already high).

Full source code (Google Colab-ready)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
# =========================================================
# Optimal Allocation of a Limited Cancer-Screening Budget
# =========================================================
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D # enables 3D projection
from scipy.optimize import minimize
import time

# ---------------------------------------------------------
# 1. Problem data
# ---------------------------------------------------------
programs = ["Breast", "Colorectal", "Lung", "Cervical"]
a = np.array([120.0, 90.0, 150.0, 60.0]) # saturation level (max detectable cases)
b = np.array([25.0, 15.0, 40.0, 10.0]) # diminishing-returns scale (in $10,000 units)

TOTAL_BUDGET = 100.0 # $10,000 units -> $1,000,000 total
n = len(programs)

def benefit(x, a_i, b_i):
return a_i * (1.0 - np.exp(-x / b_i))

def marginal_benefit(x, a_i, b_i):
return (a_i / b_i) * np.exp(-x / b_i)

def neg_total_benefit(x, a_arr, b_arr):
return -np.sum(benefit(x, a_arr, b_arr))

# ---------------------------------------------------------
# 2. Baseline solver: general-purpose optimizer (SLSQP)
# ---------------------------------------------------------
x0 = np.full(n, TOTAL_BUDGET / n)
bounds = [(0, TOTAL_BUDGET) for _ in range(n)]
constraints = {"type": "eq", "fun": lambda x: np.sum(x) - TOTAL_BUDGET}

t0 = time.time()
result = minimize(neg_total_benefit, x0, args=(a, b), method="SLSQP",
bounds=bounds, constraints=constraints)
t1 = time.time()
x_opt_slsqp = result.x
time_slsqp_single = t1 - t0

print("=== SLSQP baseline result ===")
summary = pd.DataFrame({
"Program": programs,
"Budget ($10k)": np.round(x_opt_slsqp, 2),
"Detected cases": np.round(benefit(x_opt_slsqp, a, b), 2)
})
print(summary.to_string(index=False))
print(f"Total detected cases: {benefit(x_opt_slsqp, a, b).sum():.2f}")
print(f"Solve time: {time_slsqp_single*1000:.2f} ms\n")

# ---------------------------------------------------------
# 3. Fast solver: vectorized water-filling (closed-form + bisection)
# ---------------------------------------------------------
def water_filling_vectorized(a_arr, b_arr, budgets, max_iter=100):
"""
Solves the allocation problem for MANY total budgets at once,
using the marginal-equality (water-filling) condition and
bisection on the multiplier lambda. Fully vectorized: no
Python-level loop over budgets or programs.
"""
budgets = np.asarray(budgets, dtype=float)
m = len(budgets)
lo = np.zeros(m)
hi = np.full(m, np.max(a_arr / b_arr))

for _ in range(max_iter):
lam = (lo + hi) / 2.0 # shape (m,)
ratio = np.clip(lam[:, None] * b_arr[None, :] / a_arr[None, :],
1e-300, None)
x = -b_arr[None, :] * np.log(ratio) # shape (m, n)
x = np.clip(x, 0.0, None)
over_budget = x.sum(axis=1) > budgets
lo = np.where(over_budget, lam, lo)
hi = np.where(~over_budget, lam, hi)

lam = (lo + hi) / 2.0
ratio = np.clip(lam[:, None] * b_arr[None, :] / a_arr[None, :], 1e-300, None)
x = -b_arr[None, :] * np.log(ratio)
x = np.clip(x, 0.0, None)
return x # shape (m, n)

# Verify it matches SLSQP for the baseline budget
x_opt_wf = water_filling_vectorized(a, b, [TOTAL_BUDGET])[0]
print("=== Water-filling result (should match SLSQP) ===")
print(np.round(x_opt_wf, 2))

# ---------------------------------------------------------
# 4. Speed comparison over a budget sweep
# ---------------------------------------------------------
budgets_sweep = np.linspace(1.0, 300.0, 500)

# Slow way: call SLSQP once per budget value (only 60 points, to keep runtime sane)
sample_budgets = np.linspace(1.0, 300.0, 60)
t0 = time.time()
for bud in sample_budgets:
x0_s = np.full(n, bud / n)
bnds = [(0, bud) for _ in range(n)]
cons = {"type": "eq", "fun": lambda x, bb=bud: np.sum(x) - bb}
minimize(neg_total_benefit, x0_s, args=(a, b), method="SLSQP",
bounds=bnds, constraints=cons)
t1 = time.time()
time_slsqp_sweep = t1 - t0

# Fast way: one vectorized call for all 500 budgets
t0 = time.time()
X_sweep = water_filling_vectorized(a, b, budgets_sweep)
t1 = time.time()
time_wf_sweep = t1 - t0

print(f"\nSLSQP: {len(sample_budgets)} budgets solved in {time_slsqp_sweep:.4f} s "
f"({time_slsqp_sweep/len(sample_budgets)*1000:.2f} ms/budget)")
print(f"Vectorized water-filling: {len(budgets_sweep)} budgets solved in "
f"{time_wf_sweep:.4f} s ({time_wf_sweep/len(budgets_sweep)*1000:.4f} ms/budget)")

total_benefit_sweep = np.array([
benefit(X_sweep[k], a, b).sum() for k in range(len(budgets_sweep))
])

# ---------------------------------------------------------
# 5. Graph 1: Optimal allocation (bar chart)
# ---------------------------------------------------------
fig1, ax1 = plt.subplots(figsize=(7, 5))
colors = ["#4C72B0", "#DD8452", "#55A868", "#C44E52"]
ax1.bar(programs, x_opt_wf, color=colors)
for i, v in enumerate(x_opt_wf):
ax1.text(i, v + 1, f"{v:.1f}", ha="center", fontweight="bold")
ax1.set_ylabel("Allocated budget ($10,000 units)")
ax1.set_title(f"Optimal Budget Allocation (Total = ${TOTAL_BUDGET*10000:,.0f})")
plt.tight_layout()
plt.show()

# ---------------------------------------------------------
# 6. Graph 2: Total detected cases vs. total budget
# ---------------------------------------------------------
fig2, ax2 = plt.subplots(figsize=(7, 5))
ax2.plot(budgets_sweep * 10000, total_benefit_sweep, color="#4C72B0", linewidth=2)
ax2.axvline(TOTAL_BUDGET * 10000, color="gray", linestyle="--", alpha=0.7)
ax2.set_xlabel("Total budget ($)")
ax2.set_ylabel("Total detected cases (optimal allocation)")
ax2.set_title("Diminishing Returns of the Overall Screening Budget")
plt.tight_layout()
plt.show()

# ---------------------------------------------------------
# 7. Graph 3: Marginal benefit curves (illustrates water-filling)
# ---------------------------------------------------------
x_range = np.linspace(0.01, 60, 300)
lam_opt = marginal_benefit(x_opt_wf[0], a[0], b[0]) # marginal value at optimum

fig3, ax3 = plt.subplots(figsize=(7, 5))
for i, p in enumerate(programs):
ax3.plot(x_range, marginal_benefit(x_range, a[i], b[i]),
label=p, color=colors[i])
ax3.axvline(x_opt_wf[i], color=colors[i], linestyle=":", alpha=0.6)
ax3.axhline(lam_opt, color="black", linestyle="--", label="Optimal marginal value (λ)")
ax3.set_xlabel("Budget allocated to program ($10,000 units)")
ax3.set_ylabel("Marginal benefit (extra cases per $10,000)")
ax3.set_title("Marginal Benefit Curves and the Water-Filling Condition")
ax3.legend()
plt.tight_layout()
plt.show()

# ---------------------------------------------------------
# 8. Graph 4: 3D benefit landscape (Lung vs. Breast allocation)
# ---------------------------------------------------------
grid_n = 60
xl_range = np.linspace(0, TOTAL_BUDGET, grid_n) # Lung budget
xb_range = np.linspace(0, TOTAL_BUDGET, grid_n) # Breast budget
XL, XB = np.meshgrid(xl_range, xb_range)

remaining = TOTAL_BUDGET - XL - XB
infeasible = remaining < 0
remaining_clipped = np.where(infeasible, 0.0, remaining)

# Split the remaining budget between Colorectal and Cervical
# proportionally to their saturation potential
w_col = a[1] / (a[1] + a[3])
XC = remaining_clipped * w_col
XV = remaining_clipped * (1 - w_col)

Z = (benefit(XB, a[0], b[0]) + benefit(XC, a[1], b[1]) +
benefit(XL, a[2], b[2]) + benefit(XV, a[3], b[3]))
Z[infeasible] = np.nan

fig4 = plt.figure(figsize=(9, 7))
ax4 = fig4.add_subplot(111, projection="3d")
surf = ax4.plot_surface(XL, XB, Z, cmap="viridis", edgecolor="none", alpha=0.9)
ax4.scatter([x_opt_wf[2]], [x_opt_wf[0]],
[benefit(x_opt_wf, a, b).sum()],
color="red", s=60, label="Optimal point")
ax4.set_xlabel("Lung budget ($10k)")
ax4.set_ylabel("Breast budget ($10k)")
ax4.set_zlabel("Total detected cases")
ax4.set_title("Benefit Landscape: Lung vs. Breast Allocation\n(Colorectal & Cervical fill remaining budget)")
fig4.colorbar(surf, shrink=0.5, aspect=10, label="Total detected cases")
plt.tight_layout()
plt.show()
=== SLSQP baseline result ===
   Program  Budget ($10k)  Detected cases
    Breast          28.97           82.34
Colorectal          20.73           67.40
      Lung          36.48           89.74
  Cervical          13.82           44.93
Total detected cases: 284.42
Solve time: 26.39 ms

=== Water-filling result (should match SLSQP) ===
[28.97 20.73 36.48 13.82]

SLSQP: 60 budgets solved in 0.4973 s (8.29 ms/budget)
Vectorized water-filling: 500 budgets solved in 0.0085 s (0.0170 ms/budget)

Code walkthrough

Section 1–2 (data and baseline solver). The four programs’ parameters ($a_i$, $b_i$) are defined as NumPy arrays. benefit() implements the saturation curve; neg_total_benefit() is its negative, since scipy.optimize.minimize only minimizes, never maximizes. We hand this to SLSQP (Sequential Least Squares Programming), a general-purpose constrained optimizer that can handle the equality constraint $\sum x_i = 100$ and the bounds $x_i \ge 0$. This is our “ground truth” baseline — accurate, but relatively slow because at every iteration it computes gradients and refines the solution numerically.

Section 3 (the fast, specialized solver). Instead of treating this as a generic nonlinear program, we exploit the mathematical structure. Because the benefit functions are concave, we know the optimum occurs when every program’s marginal benefit equals the same constant $\lambda$. water_filling_vectorized() searches for that $\lambda$ using bisection: it guesses a value of $\lambda$, computes how much budget each program would want at that marginal value, checks whether the total exceeds the given budget, and narrows the search range accordingly. It does this for a whole array of different total budgets simultaneously, using NumPy broadcasting (lam[:, None] * b_arr[None, :]), so there is no Python-level loop over budgets or over programs — only a fixed 100-iteration bisection loop that runs entirely in optimized NumPy C code.

Section 4 (speed comparison). We sweep 500 different total-budget scenarios (from $10,000 to $3,000,000) to see how the optimal allocation and total benefit change as the budget grows. Running SLSQP once per scenario carries real overhead (numerical differentiation, constraint handling, convergence checks), so we only ran it for 60 sample points to keep the runtime reasonable. The vectorized water-filling method solves all 500 scenarios in the same or less time that SLSQP needs for a single budget, because the entire sweep is one batch of array operations instead of hundreds of independent optimizer calls.

Sections 5–8 (visualization). These build the four charts described below.

Understanding the graphs

Graph 1 — Optimal allocation (bar chart). Shows exactly how the $1,000,000 budget should be split across the four programs at the optimum. Programs with a high ceiling ($a_i$) but that also converge slowly (large $b_i$), like Lung, tend to receive a larger share, while cheap-to-saturate programs like Cervical receive comparatively less once their curve has flattened out.

Graph 2 — Total detected cases vs. total budget. This line traces the overall diminishing-returns curve of the entire screening system, not just one program. Early dollars are extremely productive; as the total budget grows past a certain point, each additional $100,000 buys noticeably fewer additional detected cases. This is exactly the kind of curve a health ministry would want when arguing for (or against) additional funding.

Graph 3 — Marginal benefit curves. This is the most conceptually important chart. Each colored curve shows how much extra benefit one more $10,000 buys a given program, as a function of how much it has already received. The dashed horizontal line marks the common marginal value $\lambda$ at the optimum, and the dotted vertical lines show where each program’s curve crosses that line. Visually, this proves the water-filling principle: money keeps flowing into whichever program currently has the highest marginal payoff, until all four curves cross the same horizontal line — at that point, moving a dollar from one program to another can no longer improve the total.

Graph 4 — 3D benefit landscape. Here we fix the total budget at $1,000,000 and let the Lung and Breast allocations vary freely across the horizontal plane, while the leftover money automatically fills Colorectal and Cervical proportionally to their potential. The height of the surface is the total number of detected cases. Because all four benefit functions are concave, the whole landscape is a smooth dome — there’s a single peak, marked with a red dot, and no isolated local optima to get stuck in. This is a nice visual confirmation that a simple hill-climbing or water-filling approach is guaranteed to find the true best allocation.

Takeaways

This is a small model, but the pattern generalizes directly to real health-policy decisions: whenever benefits from different investments show diminishing returns, the key question is never “how much can I afford for program X” in isolation — it’s “where is a marginal dollar currently doing the most good,” across all competing programs at once. The water-filling algorithm formalizes that intuition and, as a bonus, turns out to be dramatically faster than a generic optimizer when you need to explore many what-if budget scenarios.

Hospital Radiation Therapy Equipment Scheduling

Minimizing Patient Waiting Time

Radiation oncology departments are chronically capacity-constrained. A hospital typically owns only two or three linear accelerators (LINACs), yet must serve dozens of patients per day, each with a different treatment duration and a different level of clinical urgency (a pediatric emergency case cannot wait behind a routine follow-up). The scheduling question — which patient goes on which machine, and in what order — directly determines how long people sit in the waiting room before they can start treatment.

This is a textbook instance of the unrelated/identical parallel-machine scheduling problem with weighted waiting time minimization, which is NP-hard in general. In this article we build a concrete example with 16 patients and 3 LINAC machines, solve it with a metaheuristic optimizer, and compare it against the naive first-come-first-served (FCFS) policy that many clinics still use informally.

Problem Formulation

Let $P = {1, \dots, n}$ be the set of patients and $M = {1, \dots, m}$ the set of LINAC machines. Each patient $i$ has:

  • a treatment session duration $p_i$ (minutes),
  • an arrival / ready time $r_i$ (minutes after the department opens),
  • a clinical urgency weight $w_i$.

We must decide a start time $S_i \ge r_i$ and machine assignment for every patient such that no two patients overlap on the same machine. The objective is:

$$
\min \sum_{i \in P} w_i \left(S_i - r_i\right)
$$

subject to, for every pair of patients $i,j$ sharing a machine:

$$
S_i + p_i \le S_j \quad \text{or} \quad S_j + p_j \le S_i
$$

This is a generalization of the classical $1 ,|, \sum w_i C_i$ scheduling problem, and with parallel machines it becomes combinatorially explosive very quickly ($m^n$ machine assignments alone), so exact enumeration is impractical even for modest $n$.

Solution Approach: Random-Key Encoding + Differential Evolution

Rather than optimizing discrete assignments directly, we use a well-known trick from genetic-algorithm literature called random-key encoding:

  1. Every patient $i$ is given a continuous “priority key” $k_i \in [0,1]$.
  2. A decoder sorts patients by key value (highest priority first) and greedily assigns each one, in that order, to whichever machine becomes free earliest — but never before the patient’s own arrival time $r_i$.
  3. This decoding step always produces a 100% feasible schedule, no matter what the key vector looks like.

Because the decoder guarantees feasibility, the search problem reduces to optimizing a continuous vector $\mathbf{k} \in [0,1]^n$, which is exactly what scipy.optimize.differential_evolution (DE) is built for. DE explores the key space with mutation and crossover across a population of candidate vectors, and the decoder translates each candidate into a real, evaluable schedule.

Compared to a full mixed-integer programming formulation, this approach avoids exponential blow-up, decodes in $O(n \log n)$ time, and consistently finds strong (though not provably optimal) solutions within a few hundred generations — well within Google Colab’s free-tier time budget.

Full Source Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
"""
Hospital Radiation Therapy (LINAC) Scheduling
Minimizing Weighted Patient Waiting Time
Single-file script for Google Colaboratory
"""

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.optimize import differential_evolution
import matplotlib.colors as mcolors

# ----------------------------------------------------------------------
# 0. Dark theme setup
# ----------------------------------------------------------------------
plt.rcParams.update({
"figure.facecolor": "#0d1117",
"axes.facecolor": "#0d1117",
"axes.edgecolor": "#8b949e",
"axes.labelcolor": "#e6edf3",
"text.color": "#e6edf3",
"xtick.color": "#c9d1d9",
"ytick.color": "#c9d1d9",
"grid.color": "#30363d",
"font.size": 11,
})

RNG_SEED = 42
rng = np.random.default_rng(RNG_SEED)

# ----------------------------------------------------------------------
# 1. Problem definition
# ----------------------------------------------------------------------
N_PATIENTS = 16
N_MACHINES = 3 # number of LINAC units available

durations = rng.integers(15, 46, size=N_PATIENTS).astype(float) # minutes
arrivals = rng.integers(0, 121, size=N_PATIENTS).astype(float) # minutes after opening
weights = rng.choice([1.0, 2.0, 3.0, 5.0], size=N_PATIENTS,
p=[0.4, 0.3, 0.2, 0.1]) # clinical urgency

patient_ids = np.array([f"P{i+1:02d}" for i in range(N_PATIENTS)])

# ----------------------------------------------------------------------
# 2. Decoder: priority key vector -> feasible parallel-machine schedule
# ----------------------------------------------------------------------
def decode_schedule(priority_keys, durations, arrivals, n_machines):
order = np.argsort(-priority_keys)
n = len(durations)
machine_free = np.zeros(n_machines)
start_times = np.empty(n)
machine_of = np.empty(n, dtype=int)

for idx in order:
m = int(np.argmin(machine_free))
start = machine_free[m] if machine_free[m] > arrivals[idx] else arrivals[idx]
start_times[idx] = start
machine_of[idx] = m
machine_free[m] = start + durations[idx]

return start_times, machine_of

def fcfs_schedule(durations, arrivals, n_machines):
order = np.argsort(arrivals)
n = len(durations)
machine_free = np.zeros(n_machines)
start_times = np.empty(n)
machine_of = np.empty(n, dtype=int)

for idx in order:
m = int(np.argmin(machine_free))
start = machine_free[m] if machine_free[m] > arrivals[idx] else arrivals[idx]
start_times[idx] = start
machine_of[idx] = m
machine_free[m] = start + durations[idx]

return start_times, machine_of

# ----------------------------------------------------------------------
# 3. Objective function: total weighted waiting time
# ----------------------------------------------------------------------
def weighted_wait(start_times, arrivals, weights):
return float(np.sum(weights * (start_times - arrivals)))

def objective(priority_keys):
start_times, _ = decode_schedule(priority_keys, durations, arrivals, N_MACHINES)
return weighted_wait(start_times, arrivals, weights)

# ----------------------------------------------------------------------
# 4. Optimize with Differential Evolution
# ----------------------------------------------------------------------
convergence_history = []

def de_callback(*args):
# Works across scipy versions: old-style callback(xk, convergence)
# and new-style callback(intermediate_result)
if len(args) == 1 and hasattr(args[0], "x"):
xk = args[0].x
else:
xk = args[0]
convergence_history.append(objective(xk))

bounds = [(0.0, 1.0)] * N_PATIENTS

result = differential_evolution(
objective,
bounds,
seed=RNG_SEED,
maxiter=200,
popsize=15,
mutation=(0.4, 1.0),
recombination=0.7,
tol=1e-8,
polish=True,
updating="deferred",
workers=1,
callback=de_callback,
)

best_keys = result.x
opt_start, opt_machine = decode_schedule(best_keys, durations, arrivals, N_MACHINES)
opt_wait_total = weighted_wait(opt_start, arrivals, weights)

fcfs_start, fcfs_machine = fcfs_schedule(durations, arrivals, N_MACHINES)
fcfs_wait_total = weighted_wait(fcfs_start, arrivals, weights)

improvement_pct = 100.0 * (fcfs_wait_total - opt_wait_total) / fcfs_wait_total

print(f"FCFS total weighted wait : {fcfs_wait_total:9.1f} patient-weighted-minutes")
print(f"Optimized total weighted wait : {opt_wait_total:9.1f} patient-weighted-minutes")
print(f"Improvement : {improvement_pct:6.2f} %")
print()
print(f"{'ID':4s}{'Machine':>9s}{'Arrival':>10s}{'Start':>8s}{'Wait':>7s}{'Weight':>8s}")
for i in np.argsort(opt_start):
print(f"{patient_ids[i]:4s}{opt_machine[i]+1:9d}{arrivals[i]:10.0f}"
f"{opt_start[i]:8.0f}{opt_start[i]-arrivals[i]:7.0f}{weights[i]:8.1f}")

# ----------------------------------------------------------------------
# 5. Graph 1: 3D Gantt chart of the optimized schedule
# ----------------------------------------------------------------------
fig1 = plt.figure(figsize=(11, 7))
ax1 = fig1.add_subplot(111, projection="3d")
ax1.set_facecolor("#0d1117")

cmap = plt.get_cmap("plasma")
norm = mcolors.Normalize(vmin=weights.min(), vmax=weights.max())

for i in range(N_PATIENTS):
x = opt_machine[i]
y = opt_start[i]
z0 = 0
dx, dy, dz = 0.6, durations[i] * 0.9, 1.0
color = cmap(norm(weights[i]))
ax1.bar3d(x - dx / 2, y, z0, dx, dy, dz, color=color, edgecolor="#0d1117",
shade=True, alpha=0.95)

ax1.set_xlabel("LINAC machine")
ax1.set_ylabel("Time (minutes from opening)")
ax1.set_zlabel("")
ax1.set_xticks(range(N_MACHINES))
ax1.set_xticklabels([f"LINAC-{m+1}" for m in range(N_MACHINES)])
ax1.set_zticks([])
ax1.set_title("Optimized Treatment Schedule (bar length = session duration, color = urgency)",
color="#e6edf3")
mappable = plt.cm.ScalarMappable(norm=norm, cmap=cmap)
mappable.set_array([])
cbar = fig1.colorbar(mappable, ax=ax1, shrink=0.6, pad=0.1)
cbar.set_label("Clinical priority weight")
ax1.view_init(elev=22, azim=-60)
plt.tight_layout()
plt.show()

# ----------------------------------------------------------------------
# 6. Graph 2: per-patient waiting time, FCFS vs optimized
# ----------------------------------------------------------------------
fig2, ax2 = plt.subplots(figsize=(12, 5))
x_pos = np.arange(N_PATIENTS)
bar_w = 0.38

fcfs_wait = fcfs_start - arrivals
opt_wait = opt_start - arrivals

ax2.bar(x_pos - bar_w / 2, fcfs_wait, width=bar_w, label="FCFS baseline",
color="#58a6ff", alpha=0.85)
ax2.bar(x_pos + bar_w / 2, opt_wait, width=bar_w, label="DE-optimized",
color="#f78166", alpha=0.9)
ax2.set_xticks(x_pos)
ax2.set_xticklabels(patient_ids, rotation=90)
ax2.set_ylabel("Waiting time (minutes)")
ax2.set_title("Per-Patient Waiting Time: Baseline vs Optimized Schedule", color="#e6edf3")
ax2.legend(facecolor="#161b22", edgecolor="#30363d")
ax2.grid(axis="y", alpha=0.3)
plt.tight_layout()
plt.show()

# ----------------------------------------------------------------------
# 7. Graph 3: DE convergence history
# ----------------------------------------------------------------------
fig3, ax3 = plt.subplots(figsize=(10, 5))
ax3.plot(convergence_history, color="#3fb950", linewidth=2)
ax3.set_xlabel("Generation")
ax3.set_ylabel("Best weighted waiting time (minutes)")
ax3.set_title("Differential Evolution Convergence", color="#e6edf3")
ax3.grid(alpha=0.3)
plt.tight_layout()
plt.show()

# ----------------------------------------------------------------------
# 8. Graph 4: 3D sensitivity surface (machines x patient load)
# ----------------------------------------------------------------------
def greedy_priority_ratio_schedule(durations, arrivals, weights, n_machines):
ratio = weights / durations
order = np.argsort(-ratio)
n = len(durations)
machine_free = np.zeros(n_machines)
start_times = np.empty(n)
for idx in order:
m = int(np.argmin(machine_free))
start = machine_free[m] if machine_free[m] > arrivals[idx] else arrivals[idx]
start_times[idx] = start
machine_free[m] = start + durations[idx]
return start_times

machine_range = np.arange(2, 7) # 2 .. 6 LINAC units
load_factors = np.linspace(0.6, 1.6, 11) # scales average patient duration

M_grid, L_grid = np.meshgrid(machine_range, load_factors)
Z_grid = np.zeros_like(M_grid, dtype=float)

for i in range(M_grid.shape[0]):
for j in range(M_grid.shape[1]):
scaled_durations = durations * L_grid[i, j]
s = greedy_priority_ratio_schedule(scaled_durations, arrivals, weights,
int(M_grid[i, j]))
Z_grid[i, j] = weighted_wait(s, arrivals, weights)

fig4 = plt.figure(figsize=(11, 7))
ax4 = fig4.add_subplot(111, projection="3d")
ax4.set_facecolor("#0d1117")
surf = ax4.plot_surface(M_grid, L_grid, Z_grid, cmap="viridis",
edgecolor="#0d1117", linewidth=0.3, antialiased=True)
ax4.set_xlabel("Number of LINAC machines")
ax4.set_ylabel("Patient duration load factor")
ax4.set_zlabel("Total weighted wait (min)")
ax4.set_title("Sensitivity of Total Weighted Waiting Time\nto Machine Count and Patient Load",
color="#e6edf3")
fig4.colorbar(surf, ax=ax4, shrink=0.6, pad=0.1, label="Weighted wait (min)")
ax4.view_init(elev=25, azim=-50)
plt.tight_layout()
plt.show()
FCFS total weighted wait      :     987.0 patient-weighted-minutes
Optimized total weighted wait :     570.0 patient-weighted-minutes
Improvement                   :  42.25 %

ID    Machine   Arrival   Start   Wait  Weight
P02         1        15      15      0     1.0
P07         3        22      22      0     5.0
P16         2        27      27      0     5.0
P11         3        48      48      0     2.0
P15         1        54      54      0     3.0
P05         2        60      66      6     3.0
P01         3        62      79     17     2.0
P04         1        54      91     37     2.0
P09         2        94      94      0     3.0
P10         3        77      96     19     1.0
P03         3       101     113     12     3.0
P08         2       112     115      3     3.0
P14         1        53     119     66     2.0
P12         3        99     148     49     1.0
P06         2        44     151    107     1.0
P13         1        65     157     92     1.0

Code Walkthrough

Section 0–1 (setup and problem data): We fix a random seed so the example is fully reproducible, then generate 16 synthetic patients with randomized session durations (15–45 minutes), arrival times spread across a two-hour window, and urgency weights drawn from a skewed distribution — most patients are routine (weight 1), a smaller number are moderately urgent, and a rare few are high-priority cases (weight 5). Three LINAC machines are available, matching a typical mid-sized radiation oncology department.

Section 2 (decoder): decode_schedule is the heart of the encoding trick. Given any vector of priority keys, it processes patients from highest key to lowest, and at each step assigns the current patient to whichever machine becomes idle earliest (np.argmin(machine_free)), while respecting that treatment cannot start before the patient physically arrives. This guarantees a valid, non-overlapping schedule for any input vector, which is exactly the property differential evolution needs to search freely without generating invalid candidates. fcfs_schedule is the same list-scheduling logic but with a fixed processing order (arrival time), representing the naive baseline most clinics fall back to without formal optimization.

Section 3–4 (objective and optimizer): The objective sums each patient’s waiting time weighted by clinical urgency. differential_evolution searches the 16-dimensional key space; because the decoder is $O(n \log n)$ and trivial to evaluate, even a population of 15 candidates over 200 generations evaluates in well under a second per generation, keeping total runtime short on Colab’s CPU. updating="deferred" combined with workers=1 keeps the search deterministic and reproducible given the fixed seed, while still allowing you to switch to workers=-1 for parallel evaluation if you scale up the patient count. The de_callback function records the best objective value at every generation for the convergence plot, and is written to tolerate both older and newer SciPy callback signatures so it won’t break regardless of the SciPy version bundled with your Colab runtime.

Section 5–8 (results and visualization): After optimization, we decode the best key vector into a concrete final schedule, compute the FCFS baseline for comparison, and print a per-patient summary table together with the percentage improvement in total weighted waiting time.

Graph 1 — 3D Gantt Chart of the Optimized Schedule

This chart plots each LINAC machine along the x-axis and time-of-day along the y-axis, with each patient’s treatment session rendered as a 3D bar whose length corresponds to session duration. Bar color encodes clinical urgency (brighter/warmer colors from the plasma colormap indicate higher-priority patients). This view lets you visually confirm two things at a glance: that no two bars overlap in time on the same machine (feasibility), and that the optimizer tends to slot high-urgency patients earlier and pack sessions tightly to avoid idle machine time.

Graph 2 — Per-Patient Waiting Time: Baseline vs Optimized

A direct side-by-side bar comparison of how long each individual patient waits under FCFS versus the optimized schedule. This is often the most persuasive chart for hospital administrators, since it shows concretely which patients benefit — typically the high-weight, urgent cases see their waiting time drop sharply, sometimes at the cost of slightly longer waits for low-priority routine patients, which is the intended clinical trade-off.

Graph 3 — Differential Evolution Convergence

This line plot tracks the best (lowest) weighted waiting time found at each generation of the optimizer. A healthy convergence curve should drop quickly in the first few dozen generations and then flatten out as the population converges toward a strong local optimum. If your curve is still steeply decreasing at generation 200, that’s a signal to increase maxiter; if it flattens out well before generation 200, you can safely reduce maxiter to shorten runtime without sacrificing solution quality.

Graph 4 — Sensitivity Surface: Machine Count vs Patient Load

This is a genuine capacity-planning tool. The surface plots total weighted waiting time as a function of two variables a hospital administrator actually controls or forecasts: the number of LINAC machines installed, and a “load factor” scaling how long treatment sessions run on average (a proxy for case-mix complexity, e.g. more proton therapy or IMRT cases with longer setup times). Reading across the surface at a fixed load factor shows the diminishing-returns curve of adding machines — useful for justifying (or questioning) capital equipment purchases. Reading along a fixed machine count shows how sensitive the department is to creeping session durations, which is valuable for staffing and throughput planning.

Conclusion

Framing LINAC scheduling as a weighted parallel-machine problem and solving it with a random-key differential evolution approach turns an informal, often first-come-first-served process into a quantifiable optimization with measurable improvement. The same decoder-plus-metaheuristic pattern generalizes well beyond radiation oncology — it applies equally to CT/MRI scanner scheduling, operating room allocation, or any other constrained clinical resource where patient urgency and equipment availability must be balanced against waiting time.

Optimizing Clinical Trial Design

How Many Patients, and What Dose?

Every clinical trial faces two deceptively simple questions before a single patient is enrolled: how many subjects do we need? and what dose should we test? Get the sample size wrong, and the trial either wastes resources on unnecessary patients or fails to detect a real treatment effect. Get the dose wrong, and patients are exposed to either ineffective or dangerously toxic treatment levels.

In this article, we tackle both problems computationally. First, we determine the optimal sample size for a two-arm superiority trial using statistical power analysis, verified with a vectorized Monte Carlo simulation. Second, we simulate a Bayesian dose-escalation trial using the Continual Reassessment Method (CRM) to efficiently find the Maximum Tolerated Dose (MTD) while minimizing patient exposure to toxicity.

The Statistical Foundation

Sample Size

For a two-arm trial comparing a treatment mean against a control mean, the number of patients needed per arm to detect a true difference $\Delta$ with a given significance level $\alpha$ and power $1-\beta$ is:

$$
n = \frac{2(z_{\alpha/2} + z_{\beta})^2 \sigma^2}{\Delta^2}
$$

where $\sigma$ is the assumed standard deviation of the outcome, $z_{\alpha/2}$ is the critical value for the significance level, and $z_{\beta}$ is the critical value corresponding to the desired power. Intuitively: smaller effect sizes, larger variance, or stricter significance/power requirements all demand more patients.

Dose Finding

For dose-escalation trials, we model the probability of a dose-limiting toxicity (DLT) at dose $d_i$ using a one-parameter power model:

$$
p(d_i) = \psi_i^{\beta}
$$

where $\psi_i$ is a prior “skeleton” toxicity guess for dose level $i$, and $\beta$ is a parameter updated in real time as patient outcomes accumulate. The trial seeks the dose whose estimated toxicity is closest to a pre-specified target, typically $p(d_i) \approx 0.30$.

The Code

Below is the complete, self-contained script. It runs top to bottom in Google Colaboratory with no additional setup.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
# ===================================================
# Clinical Trial Design Optimization
# Part A: Sample Size Determination (Two-Arm Superiority Trial)
# Part B: Dose-Finding via Bayesian Continual Reassessment Method (CRM)
# ===================================================

import numpy as np
import scipy.stats as stats
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import time

np.random.seed(42)

# ---------------------------------------------------
# PART A: SAMPLE SIZE DETERMINATION
# ---------------------------------------------------

def required_sample_size(effect_size, sigma, alpha=0.05, power=0.8):
"""
Analytic sample size per arm for a two-sample (two-sided) test
comparing means, using the normal approximation.
"""
z_alpha = stats.norm.ppf(1 - alpha / 2)
z_beta = stats.norm.ppf(power)
n = 2 * ((z_alpha + z_beta) ** 2) * (sigma ** 2) / (effect_size ** 2)
return int(np.ceil(n))

def analytic_power(n, effect_size, sigma, alpha=0.05):
"""
Closed-form statistical power for a two-sample z-test, used to
quickly evaluate power across many (n, effect_size) combinations.
"""
z_alpha = stats.norm.ppf(1 - alpha / 2)
se = sigma * np.sqrt(2.0 / n)
ncp = effect_size / se
power = 1 - stats.norm.cdf(z_alpha - ncp) + stats.norm.cdf(-z_alpha - ncp)
return power

def monte_carlo_power(n, effect_size, sigma, alpha=0.05, n_sim=20000):
"""
Vectorized Monte Carlo estimate of statistical power.
Simulates n_sim trials simultaneously using numpy array operations
(no Python-level loop over simulations), which is essential for speed.
"""
control = np.random.normal(loc=0.0, scale=sigma, size=(n_sim, n))
treatment = np.random.normal(loc=effect_size, scale=sigma, size=(n_sim, n))

mean_diff = treatment.mean(axis=1) - control.mean(axis=1)
pooled_var = (treatment.var(axis=1, ddof=1) + control.var(axis=1, ddof=1)) / 2
se = np.sqrt(2 * pooled_var / n)
t_stat = mean_diff / se

df = 2 * (n - 1)
t_crit = stats.t.ppf(1 - alpha / 2, df)
reject = np.abs(t_stat) > t_crit
return reject.mean()

# --- Baseline design parameters ---
sigma = 1.0
alpha = 0.05
target_power = 0.8
effect_size = 0.5

n_needed = required_sample_size(effect_size, sigma, alpha, target_power)
print(f"Analytic required sample size per arm: {n_needed}")

start = time.time()
mc_power = monte_carlo_power(n_needed, effect_size, sigma, alpha, n_sim=20000)
elapsed = time.time() - start
print(f"Monte Carlo estimated power at n={n_needed}: {mc_power:.3f} (simulated in {elapsed:.2f}s)")

# --- Power curve: power vs n for several effect sizes ---
n_range = np.arange(5, 201, 5)
effect_sizes_list = [0.2, 0.35, 0.5, 0.8]

plt.figure(figsize=(8, 5))
for es in effect_sizes_list:
powers = analytic_power(n_range, es, sigma, alpha)
plt.plot(n_range, powers, label=f"effect size = {es}")
plt.axhline(0.8, color="gray", linestyle="--", linewidth=1, label="target power = 0.8")
plt.xlabel("Sample size per arm (n)")
plt.ylabel("Statistical power")
plt.title("Power vs Sample Size for Different Effect Sizes")
plt.legend()
plt.grid(alpha=0.3)
plt.tight_layout()
plt.show()
# [Insert power-vs-sample-size chart here]

# --- 3D surface: power as a function of n and effect size ---
n_grid = np.arange(5, 201, 5)
es_grid = np.linspace(0.1, 1.0, 40)
N, ES = np.meshgrid(n_grid, es_grid)
POWER = analytic_power(N, ES, sigma, alpha)

fig = plt.figure(figsize=(9, 6))
ax = fig.add_subplot(111, projection="3d")
surf = ax.plot_surface(N, ES, POWER, cmap="viridis", edgecolor="none")
ax.set_xlabel("Sample size (n)")
ax.set_ylabel("Effect size")
ax.set_zlabel("Power")
ax.set_title("Statistical Power Surface: Sample Size vs Effect Size")
fig.colorbar(surf, shrink=0.5, aspect=10, label="Power")
plt.tight_layout()
plt.show()
# [Insert 3D power surface chart here]


# ---------------------------------------------------
# PART B: DOSE OPTIMIZATION VIA BAYESIAN CRM
# ---------------------------------------------------

dose_levels = np.array([1, 2, 3, 4, 5])
skeleton = np.array([0.05, 0.12, 0.25, 0.40, 0.55]) # prior toxicity guesses per dose

def true_toxicity_prob(dose_idx):
"""
The true (unknown to the investigator) probability of dose-limiting
toxicity at each dose level, used only to generate simulated outcomes.
"""
true_curve = np.array([0.03, 0.10, 0.30, 0.50, 0.70])
return true_curve[dose_idx]

beta_grid = np.linspace(0.1, 5.0, 400)
prior = np.ones_like(beta_grid)
prior /= prior.sum()

target_toxicity = 0.30
n_patients = 30

dose_history = []
toxicity_history = []
posterior_curve_history = []
posterior = prior.copy()

for patient in range(n_patients):
skeleton_matrix = skeleton[None, :] ** beta_grid[:, None]
mean_tox_per_dose = (posterior[:, None] * skeleton_matrix).sum(axis=0)

next_dose_idx = int(np.argmin(np.abs(mean_tox_per_dose - target_toxicity)))

true_p = true_toxicity_prob(next_dose_idx)
outcome = np.random.binomial(1, true_p)

p_dose = skeleton[next_dose_idx] ** beta_grid
likelihood = np.where(outcome == 1, p_dose, 1 - p_dose)
posterior = posterior * likelihood
posterior /= posterior.sum()

dose_history.append(next_dose_idx)
toxicity_history.append(outcome)
posterior_curve_history.append(mean_tox_per_dose.copy())

final_matrix = skeleton[None, :] ** beta_grid[:, None]
final_mean_tox = (posterior[:, None] * final_matrix).sum(axis=0)
estimated_mtd_idx = int(np.argmin(np.abs(final_mean_tox - target_toxicity)))

print(f"Estimated MTD dose level: {dose_levels[estimated_mtd_idx]} "
f"(estimated toxicity = {final_mean_tox[estimated_mtd_idx]:.3f})")
print(f"Total DLT events observed: {sum(toxicity_history)} / {n_patients} patients")

# --- Dose allocation trajectory ---
plt.figure(figsize=(8, 5))
colors = ["red" if t == 1 else "blue" for t in toxicity_history]
plt.scatter(range(1, n_patients + 1), [dose_levels[d] for d in dose_history], c=colors)
plt.plot(range(1, n_patients + 1), [dose_levels[d] for d in dose_history],
color="gray", alpha=0.4, linewidth=1)
plt.xlabel("Patient number")
plt.ylabel("Assigned dose level")
plt.title("CRM Dose Allocation Sequence (red = DLT observed, blue = no DLT)")
plt.yticks(dose_levels)
plt.grid(alpha=0.3)
plt.tight_layout()
plt.show()

# --- 3D surface: posterior toxicity estimate evolution over patients ---
posterior_curve_history = np.array(posterior_curve_history)
Patients_axis = np.arange(1, n_patients + 1)
Dose_axis = dose_levels
P_grid, D_grid = np.meshgrid(Patients_axis, Dose_axis, indexing="ij")

fig = plt.figure(figsize=(9, 6))
ax = fig.add_subplot(111, projection="3d")
surf = ax.plot_surface(D_grid, P_grid, posterior_curve_history, cmap="plasma", edgecolor="none")
ax.set_xlabel("Dose level")
ax.set_ylabel("Patient number")
ax.set_zlabel("Estimated toxicity probability")
ax.set_title("Evolution of Posterior Toxicity Estimates During the CRM Trial")
fig.colorbar(surf, shrink=0.5, aspect=10, label="Estimated toxicity")
plt.tight_layout()
plt.show()

Results

Analytic required sample size per arm: 63
Monte Carlo estimated power at n=63: 0.796 (simulated in 0.18s)


Estimated MTD dose level: 3 (estimated toxicity = 0.314)
Total DLT events observed: 10 / 30 patients


Code Walkthrough

Part A — Sample Size

required_sample_size implements the closed-form formula directly. It converts the desired significance level and power into their corresponding standard-normal quantiles ($z_{\alpha/2}$, $z_{\beta}$) and plugs them into the sample-size equation. This gives an instant answer — no simulation needed — and is the number a biostatistician would report in a protocol.

analytic_power is the inverse direction: given a candidate $n$, it computes the expected power directly from the normal approximation. Because it’s a pure closed-form calculation, it can be evaluated over thousands of $(n, \text{effect size})$ combinations almost instantly — this is what makes the 3D surface plot feasible without a heavy simulation loop.

monte_carlo_power exists to validate the analytic formula against real random sampling, since the analytic approximation relies on asymptotic normality. Critically, this function is written to avoid a Python-level loop over simulations. Instead, it draws an entire (n_sim, n) matrix of random patients for each arm in one call, and computes all n_sim t-statistics simultaneously via NumPy’s axis=1 reductions. Simulating 20,000 full trials this way finishes in roughly a quarter of a second, whereas simulating each trial with a Python for loop would take on the order of seconds to tens of seconds. The Monte Carlo result should land close to the target power (0.8), confirming the analytic formula is trustworthy.

The 2D power curve plot shows, for several effect sizes, how power climbs toward 1.0 as sample size increases — visually confirming why weaker effects require disproportionately larger trials. The 3D surface plot extends this into a full response surface over both sample size and effect size simultaneously, making it easy to read off, for any assumed effect size, exactly how many patients are needed to cross the 80% power threshold.

Part B — Dose Finding (CRM)

The CRM section simulates a Phase I dose-escalation trial. skeleton represents the trial team’s prior belief about toxicity at each of five dose levels; true_toxicity_prob represents the actual, unknown biology used only to generate simulated patient responses — a stand-in for what would happen if this were a real trial.

beta_grid discretizes the unknown model parameter $\beta$ into 400 candidate values, each carrying a posterior probability mass (this is a “grid approximation” to full Bayesian inference — simple, numerically stable, and fast).

Inside the trial loop, for each incoming patient:

  1. skeleton_matrix computes $\psi_i^{\beta}$ for every dose $i$ and every candidate $\beta$ at once (a vectorized outer power operation — no nested loops).
  2. mean_tox_per_dose averages these predictions over the current posterior distribution of $\beta$, producing the model’s current best guess of toxicity at each dose.
  3. The patient is assigned to whichever dose’s predicted toxicity is closest to the 30% target — this is the “efficient” part of the design, since it concentrates patients near the informative, clinically relevant dose rather than wasting them on doses that are obviously too safe or too toxic.
  4. A synthetic outcome (DLT or not) is drawn from the true toxicity curve.
  5. Bayes’ rule updates the posterior over $\beta$ using the observed outcome, and the loop repeats.

After 30 simulated patients, the script reports the estimated MTD and how many toxicity events occurred in total — usually a small fraction of patients, which is the point of using an adaptive design instead of fixed dosing cohorts.

The dose allocation trajectory plot shows how the assigned dose escalates, plateaus, or retreats as evidence accumulates, with red markers flagging patients who experienced a DLT. The 3D posterior evolution surface is the most informative visual: it shows the estimated toxicity curve across all five doses at every point in the trial, letting you watch the model’s beliefs sharpen and converge toward the true dose-toxicity relationship as more patients are enrolled.

Why This Matters

Both simulations reflect a common theme in modern trial design: instead of relying on fixed, rule-of-thumb allocations (like the traditional 3+3 dose-escalation scheme, or arbitrarily “round” sample sizes), adaptive and analytically-grounded methods let trials use exactly as many resources as needed — no more, no less — while exposing fewer patients to ineffective or unsafe treatment. The sample size analysis ensures a trial is statistically capable of detecting a real effect without being wastefully oversized, and the CRM simulation shows how Bayesian updating can locate a safe, effective dose with a fraction of the patients a naive escalation scheme would require.

Optimizing Personalized Medicine

Genotype-Guided Treatment Selection with Python

Precision medicine promises to replace the “one dose fits all” paradigm with treatment decisions tailored to each patient’s genome. In practice, this means combining pharmacogenomic data — variants in drug-metabolizing enzymes, transporters, and drug targets — with dose-response models to select both the right drug and the right dose for a given patient. This is fundamentally a constrained optimization problem: maximize expected clinical benefit while keeping toxicity risk below an acceptable threshold.

In this post, we build a concrete example: a patient with a known genetic profile is being considered for one of five candidate treatments. We optimize the dose for each candidate under a toxicity ceiling, then rank the treatments to recommend the one offering the best genotype-adjusted outcome.

1. From Genotype to a Composite Pharmacogenomic Score

A patient’s genetic profile is represented by allele dosages $x_i \in {0,1,2}$ at $n$ relevant SNPs (e.g., variants in CYP2D6, CYP2C19, DPYD), each with an effect weight $w_i$ describing its influence on drug clearance. We collapse this into a single interpretable metabolizer score:

$$
g = \sigma!\left(\frac{\sum_{i=1}^{n} w_i x_i - \mu}{s}\right), \qquad \sigma(z) = \frac{1}{1+e^{-z}}
$$

where $g \in (0,1)$ ranges from poor metabolizer ($g \to 0$) to ultra-rapid metabolizer ($g \to 1$), and $\mu, s$ are calibration constants centering the score on a typical population range.

2. Genetics-Modulated Dose-Response Models

For each candidate drug $k$, efficacy and toxicity follow Hill/Emax pharmacodynamic models, but their potency parameters shift with the patient’s metabolizer score:

$$
E_k(d,g) = E_{max,k}\cdot\frac{d^{,n_k}}{EC50_k(g)^{,n_k}+d^{,n_k}}, \qquad
T_k(d,g) = T_{max,k}\cdot\frac{d^{,m_k}}{TC50_k(g)^{,m_k}+d^{,m_k}}
$$

$$
EC50_k(g) = EC50_k^0\big(1+\alpha_k(g-0.5)\big), \qquad
TC50_k(g) = TC50_k^0\big(1+\beta_k(g-0.5)\big)
$$

The coefficients $\alpha_k, \beta_k$ encode how a patient’s metabolic speed shifts the effective potency and safety margin of drug $k$ — a fast metabolizer might need a higher dose to reach the same efficacy ($\alpha_k>0$), or might clear the drug fast enough to tolerate higher doses before toxicity appears ($\beta_k>0$).

3. The Optimization Problem

For a fixed patient genotype $g^*$, we choose the dose $d$ that maximizes a clinical utility function penalizing toxicity, subject to a hard toxicity ceiling:

$$
\max_{d \in [d_{min,k},,d_{max,k}]} ; B_k(d) = E_k(d,g^*) - \lambda, T_k(d,g^*)
\quad \text{s.t.} \quad T_k(d,g^*) \le T_{allow}
$$

We solve this independently for all five candidate drugs and select the one with the highest feasible net benefit — this is the genotype-guided recommendation.

4. Full Implementation

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
# ============================================================
# Personalized Medicine: Genotype-Guided Treatment Optimization
# Single-file script for Google Colaboratory
# ============================================================

import numpy as np
from scipy.optimize import differential_evolution, NonlinearConstraint
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D # noqa: F401 (3D projection registration)

plt.style.use('dark_background')
np.random.seed(42)

# ------------------------------------------------------------
# 1. Patient genetic profile -> composite pharmacogenomic score
# ------------------------------------------------------------
snp_genotype = np.array([2, 1, 0, 2, 1, 1], dtype=float)
snp_weight = np.array([0.80, -0.50, 0.30, 0.60, -0.20, 0.40])
snp_names = ['CYP2D6*4', 'CYP2C19*2', 'CYP3A5*3', 'UGT1A1*28', 'ABCB1', 'DPYD*2A']

g_raw = np.dot(snp_weight, snp_genotype)
g_baseline, g_scale = 2.0, 2.0
g_patient = 1.0 / (1.0 + np.exp(-(g_raw - g_baseline) / g_scale))

# ------------------------------------------------------------
# 2. Candidate treatments: genetics-modulated dose-response models
# ------------------------------------------------------------
drugs = {
'Drug A': dict(Emax=0.95, EC50_0=40, alpha=0.70, n=2.0,
Tmax=0.90, TC50_0=90, beta=0.55, m=2.5, dmin=5, dmax=150),
'Drug B': dict(Emax=0.85, EC50_0=25, alpha=0.20, n=1.8,
Tmax=0.75, TC50_0=60, beta=-0.60, m=2.0, dmin=2, dmax=100),
'Drug C': dict(Emax=0.90, EC50_0=55, alpha=-0.45, n=2.2,
Tmax=0.80, TC50_0=110, beta=0.30, m=2.8, dmin=5, dmax=180),
'Drug D': dict(Emax=0.80, EC50_0=15, alpha=0.10, n=1.5,
Tmax=0.95, TC50_0=35, beta=0.80, m=3.0, dmin=1, dmax=60),
'Drug E': dict(Emax=0.97, EC50_0=70, alpha=0.35, n=2.4,
Tmax=0.70, TC50_0=140, beta=-0.25, m=2.2, dmin=8, dmax=200),
}

lam = 1.0
T_allow = 0.30


def ec50_of_g(g, p):
return p['EC50_0'] * (1.0 + p['alpha'] * (g - 0.5))


def tc50_of_g(g, p):
return p['TC50_0'] * (1.0 + p['beta'] * (g - 0.5))


def efficacy(d, g, p):
ec50 = ec50_of_g(g, p)
return p['Emax'] * d ** p['n'] / (ec50 ** p['n'] + d ** p['n'])


def toxicity(d, g, p):
tc50 = tc50_of_g(g, p)
return p['Tmax'] * d ** p['m'] / (tc50 ** p['m'] + d ** p['m'])


def net_benefit(d, g, p, lam=lam):
return efficacy(d, g, p) - lam * toxicity(d, g, p)


# ------------------------------------------------------------
# 3. Per-drug dose optimization under a toxicity ceiling
# ------------------------------------------------------------
results = {}
for name, p in drugs.items():
def neg_benefit(x, p=p):
return -net_benefit(x[0], g_patient, p)

def tox_constraint_fn(x, p=p):
return toxicity(x[0], g_patient, p)

nlc = NonlinearConstraint(tox_constraint_fn, -np.inf, T_allow)

res = differential_evolution(
neg_benefit,
bounds=[(p['dmin'], p['dmax'])],
constraints=(nlc,),
seed=42,
maxiter=400,
popsize=25,
tol=1e-8,
polish=True,
mutation=(0.5, 1.5),
recombination=0.7,
)

d_opt = res.x[0]
e_opt = efficacy(d_opt, g_patient, p)
t_opt = toxicity(d_opt, g_patient, p)
b_opt = e_opt - lam * t_opt
feasible = t_opt <= T_allow + 1e-6

results[name] = dict(dose=d_opt, efficacy=e_opt, toxicity=t_opt,
benefit=b_opt, feasible=feasible)

feasible_results = {k: v for k, v in results.items() if v['feasible']}
best_drug = max(feasible_results, key=lambda k: feasible_results[k]['benefit']) \
if feasible_results else max(results, key=lambda k: results[k]['benefit'])

print(f"Patient pharmacogenomic score g = {g_patient:.3f}")
for name, r in results.items():
flag = 'OK' if r['feasible'] else 'exceeds toxicity limit'
print(f"{name}: dose={r['dose']:6.2f} efficacy={r['efficacy']:.3f} "
f"toxicity={r['toxicity']:.3f} benefit={r['benefit']:.3f} [{flag}]")
print(f"--> Recommended treatment: {best_drug}")

# ------------------------------------------------------------
# 4. Visualization
# ------------------------------------------------------------
colors = {'Drug A': '#00e5ff', 'Drug B': '#ff6ec7', 'Drug C': '#ffd166',
'Drug D': '#8ac926', 'Drug E': '#c77dff'}

fig = plt.figure(figsize=(16, 13))
fig.patch.set_facecolor('#0d1117')

# --- 4-1. 3D benefit landscape for the recommended drug ---
ax1 = fig.add_subplot(2, 2, 1, projection='3d')
p_best = drugs[best_drug]
d_grid = np.linspace(p_best['dmin'], p_best['dmax'], 90)
g_grid = np.linspace(0.01, 0.99, 90)
D, G = np.meshgrid(d_grid, g_grid)
B = net_benefit(D, G, p_best)
surf = ax1.plot_surface(D, G, B, cmap='plasma', edgecolor='none', alpha=0.92)
ax1.scatter([results[best_drug]['dose']], [g_patient], [results[best_drug]['benefit']],
color='white', s=70, edgecolor='black', depthshade=False)
ax1.set_xlabel('Dose (mg)')
ax1.set_ylabel('Genetic score g')
ax1.set_zlabel('Net benefit')
ax1.set_title(f'Benefit landscape: {best_drug}', color='white')
fig.colorbar(surf, ax=ax1, shrink=0.55, pad=0.08)

# --- 4-2. 3D comparison of all drugs across the genetic score spectrum ---
ax2 = fig.add_subplot(2, 2, 2, projection='3d')
drug_names = list(drugs.keys())
g_bins = np.linspace(0.05, 0.95, 6)
_dx, _dy = 0.6, 0.1
for i, name in enumerate(drug_names):
p = drugs[name]
d_fixed = results[name]['dose']
heights = net_benefit(d_fixed, g_bins, p)
xs = np.full_like(g_bins, i, dtype=float)
ax2.bar3d(xs, g_bins, np.zeros_like(heights), _dx, _dy, heights,
color=colors[name], alpha=0.85, shade=True)
ax2.set_xticks(range(len(drug_names)))
ax2.set_xticklabels(drug_names, rotation=15)
ax2.set_ylabel('Genetic score g')
ax2.set_zlabel('Net benefit')
ax2.set_title('Benefit sensitivity to genotype (dose fixed per drug)', color='white')

# --- 4-3. Efficacy vs toxicity trade-off (Pareto view) at patient's genotype ---
ax3 = fig.add_subplot(2, 2, 3)
for name, p in drugs.items():
dd = np.linspace(p['dmin'], p['dmax'], 200)
ee = efficacy(dd, g_patient, p)
tt = toxicity(dd, g_patient, p)
ax3.plot(tt, ee, color=colors[name], label=name, linewidth=2)
ax3.scatter([results[name]['toxicity']], [results[name]['efficacy']],
color=colors[name], edgecolor='white', s=80, zorder=5)
ax3.axvline(T_allow, color='red', linestyle='--', linewidth=1.2, label='toxicity ceiling')
ax3.set_xlabel('Toxicity probability')
ax3.set_ylabel('Efficacy probability')
ax3.set_title('Efficacy-toxicity trade-off', color='white')
ax3.legend(fontsize=8, loc='lower right')
ax3.grid(alpha=0.2)

# --- 4-4. Ranking summary ---
ax4 = fig.add_subplot(2, 2, 4)
names_sorted = sorted(results, key=lambda k: results[k]['benefit'], reverse=True)
benefits_sorted = [results[k]['benefit'] for k in names_sorted]
bar_colors = [colors[k] if results[k]['feasible'] else '#555555' for k in names_sorted]
bars = ax4.bar(names_sorted, benefits_sorted, color=bar_colors, edgecolor='white')
for bar, name in zip(bars, names_sorted):
if not results[name]['feasible']:
ax4.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
'infeasible', ha='center', color='red', fontsize=8)
ax4.set_ylabel('Optimal net benefit')
ax4.set_title('Treatment ranking for this patient', color='white')
ax4.grid(alpha=0.2, axis='y')

plt.tight_layout()
plt.show()

Output

Patient pharmacogenomic score g = 0.562
Drug A: dose= 59.67  efficacy=0.638  toxicity=0.223  benefit=0.415  [OK]
Drug B: dose= 39.63  efficacy=0.588  toxicity=0.240  benefit=0.348  [OK]
Drug C: dose= 76.43  efficacy=0.618  toxicity=0.204  benefit=0.414  [OK]
Drug D: dose= 18.68  efficacy=0.463  toxicity=0.110  benefit=0.353  [OK]
Drug E: dose=120.93  efficacy=0.756  toxicity=0.300  benefit=0.456  [OK]
--> Recommended treatment: Drug E

5. Code Walkthrough

Genetic score construction. The snp_genotype array encodes the patient’s allele dosage at six pharmacogenomic markers, and snp_weight encodes each variant’s directional effect on metabolic activity. The dot product is passed through a logistic sigmoid, producing a bounded score $g \in (0,1)$ that is easy to interpret and easy to plug into downstream pharmacodynamic equations — this mirrors how polygenic risk scores are constructed in practice.

Dose-response functions. efficacy() and toxicity() implement the Hill-equation model described above. Both are fully vectorized with NumPy, so they accept scalars, 1D arrays, or 2D meshgrids without any code changes — this is what lets the same functions serve the optimizer and the 3D surface plot.

Constrained optimization. For each drug, we use scipy.optimize.differential_evolution, a global, gradient-free optimizer well suited to the potentially non-convex trade-off between efficacy and toxicity. The toxicity ceiling is enforced with a NonlinearConstraint object (rather than the older dict-based constraint API, which is deprecated and less numerically robust). Since the search space is one-dimensional (the dose), convergence is essentially instantaneous — no further speed-up is required here.

Selection logic. After optimizing all five drugs independently, we filter to only the feasible ones (those respecting the toxicity ceiling) and pick the drug with the highest net benefit. If no drug is feasible, we fall back to the best available option and flag it — this avoids ever silently recommending a treatment that violates safety constraints.

6. Reading the Graphs

Top-left (3D surface). This shows how net clinical benefit for the recommended drug varies jointly with dose and genetic score. The white marker sits at the patient’s actual genotype and optimal dose — the peak (or near-peak) of the ridge. The shape of this surface makes it visually clear why a fixed, non-personalized dose would be suboptimal for patients at the opposite end of the metabolizer spectrum.

Top-right (3D bar comparison). Here we take each drug’s dose (optimized specifically for our patient) and ask: how would that same dose perform across a range of hypothetical genetic scores? Drugs with flatter bar profiles are more “robust” to genotype misestimation, while drugs with steep gradients are highly genotype-sensitive — valuable information when confidence in a genetic test is imperfect.

Bottom-left (efficacy-toxicity trade-off). Each curve traces out the achievable combinations of efficacy and toxicity as dose is swept across its allowed range for a given drug, evaluated at the patient’s actual genotype. The red dashed line marks the toxicity ceiling; curves that stay well left of it before saturating in efficacy are the most attractive candidates. The dots mark each drug’s optimizer-selected operating point.

Bottom-right (ranking bar chart). A direct summary: the optimal net benefit achieved by each drug, sorted from best to worst. Gray bars (if any) indicate drugs that could not satisfy the toxicity constraint at any dose and were excluded from consideration.

7. Closing Notes

This example is a simplified, illustrative model — the pharmacodynamic parameters are synthetic rather than derived from real trial data, and any real clinical decision support system would need far more rigorous, validated dose-response models, uncertainty quantification, and regulatory oversight. But the underlying optimization structure — genotype-conditioned dose-response curves, a constrained multi-objective utility function, and a global optimizer sweeping across candidate treatments — reflects the computational backbone of real pharmacogenomics-driven decision tools, and generalizes naturally to more markers, more drugs, or multi-dose regimens.

Optimizing the Order of Cancer Treatments

Surgery, Chemotherapy, and Radiation as a Combinatorial Optimization Problem

When a patient is diagnosed with a tumor that requires multiple treatment modalities, doctors don’t just choose which treatments to use — they also have to choose the order. Should surgery come first, followed by radiation to clean up the margins? Or should chemotherapy shrink the tumor first, making surgery safer and less invasive? Each ordering carries a different balance of risk, healing time, and toxicity.

This is, at its core, a sequencing optimization problem — mathematically very close to the famous Traveling Salesman Problem (TSP), except instead of minimizing travel distance between cities, we minimize a “treatment cost” that depends on which therapy follows which.

In this article, I’ll build a concrete, solvable example with three treatments — Surgery (S), Chemotherapy (C), and Radiation (R) — model the cost of every possible ordering, find the optimal sequence with Python, visualize it (including in 3D), and then scale the problem up to nine treatment modalities to show why a smarter algorithm (dynamic programming) becomes essential as the number of choices grows.

Note: the numbers used below are illustrative, synthetic values chosen to make the optimization clear — not real clinical data. The point of this article is the algorithm, not medical advice.

Formulating the Problem

Let treatments be indexed $1, 2, \dots, n$. We define two cost components:

  • $B(i)$: the baseline cost of starting the entire treatment plan with treatment $i$ (e.g., the inherent risk of operating on an untreated tumor).
  • $C(i, j)$: the transition cost of performing treatment $j$ immediately after treatment $i$ (e.g., the risk of operating on tissue that was just irradiated).

For a treatment order $\pi = (\pi_1, \pi_2, \dots, \pi_n)$, the total cost is:

$$
\text{TotalCost}(\pi) = B(\pi_1) + \sum_{k=1}^{n-1} C(\pi_k, \pi_{k+1})
$$

We want to find the permutation $\pi^*$ that minimizes this:

$$
\pi^* = \arg\min_{\pi \in S_n} \text{TotalCost}(\pi)
$$

This is exactly the structure of the shortest Hamiltonian path problem. For small $n$, we can simply check every permutation (there are $n!$ of them). For larger $n$, we’ll use the Held-Karp dynamic programming algorithm, which solves it in:

$$
O(n^2 \cdot 2^n) \quad \text{instead of} \quad O(n!)
$$

The DP recurrence builds up the optimal cost of visiting a subset of treatments $S$ and ending at treatment $j$:

$$
dp[S][j] = \min_{i \in S \setminus {j}} \Big( dp[S \setminus {j}][i] + C(i,j) \Big), \qquad dp[{j}][j] = B(j)
$$

The Python Implementation

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
import numpy as np
import itertools
import math
import time
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# ============================================================
# 1. Core 3-treatment example: Surgery, Chemotherapy, Radiation
# ============================================================
treatments_3 = ['Surgery', 'Chemotherapy', 'Radiation']
n3 = len(treatments_3)

# B[i]: baseline cost of starting the whole plan with treatment i
B3 = np.array([5.0, 3.0, 4.0])

# C[i][j]: extra cost incurred when treatment j is performed
# immediately after treatment i (toxicity overlap, healing delay, etc.)
C3 = np.array([
[0.0, 2.0, 6.0], # from Surgery -> Chemo, Radiation
[1.5, 0.0, 2.5], # from Chemotherapy -> Surgery, Radiation
[4.0, 3.0, 0.0], # from Radiation -> Surgery, Chemo
])

def total_cost(perm, B, C):
"""Total cost of a given treatment order."""
cost = B[perm[0]]
for i in range(len(perm) - 1):
cost += C[perm[i]][perm[i + 1]]
return cost

def evaluate_all_permutations(n, B, C):
"""Full brute-force enumeration (fine for small n)."""
results = []
best_perm, best_cost = None, math.inf
for perm in itertools.permutations(range(n)):
c = total_cost(perm, B, C)
results.append((perm, c))
if c < best_cost:
best_perm, best_cost = perm, c
return best_perm, best_cost, results

best_perm3, best_cost3, all_results3 = evaluate_all_permutations(n3, B3, C3)

print("All possible treatment sequences and their total cost:")
for perm, c in all_results3:
seq_name = " -> ".join(treatments_3[i] for i in perm)
marker = " <== OPTIMAL" if perm == best_perm3 else ""
print(f"{seq_name:35s} cost = {c:5.2f}{marker}")

print("\nOptimal treatment order:",
" -> ".join(treatments_3[i] for i in best_perm3))
print("Minimum total cost:", round(best_cost3, 2))


# ============================================================
# 2. Bar chart of all 6 possible sequences
# ============================================================
labels = [" -> ".join(treatments_3[i] for i in perm) for perm, _ in all_results3]
costs = [c for _, c in all_results3]
colors = ['#ff6b6b' if perm == best_perm3 else '#4dabf7' for perm, _ in all_results3]

plt.figure(figsize=(10, 5))
plt.bar(labels, costs, color=colors)
plt.ylabel("Total treatment-plan cost")
plt.title("Total Cost for Every Possible Surgery / Chemo / Radiation Order")
plt.xticks(rotation=20, ha='right')
plt.tight_layout()
plt.show()


# ============================================================
# 3. 3D visualization of the transition-cost matrix
# ============================================================
fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, projection='3d')

xpos, ypos = np.meshgrid(np.arange(n3), np.arange(n3), indexing='ij')
xpos = xpos.flatten().astype(float)
ypos = ypos.flatten().astype(float)
zpos = np.zeros_like(xpos)
dx = dy = 0.5 * np.ones_like(xpos)
dz = C3.flatten()

colors_3d = plt.cm.viridis(dz / dz.max())
ax.bar3d(xpos, ypos, zpos, dx, dy, dz, color=colors_3d, shade=True)

ax.set_xticks(np.arange(n3) + 0.25)
ax.set_xticklabels(treatments_3)
ax.set_yticks(np.arange(n3) + 0.25)
ax.set_yticklabels(treatments_3)
ax.set_xlabel("From")
ax.set_ylabel("To")
ax.set_zlabel("Transition cost")
ax.set_title("Transition Cost Between Treatments (3D View)")
plt.tight_layout()
plt.show()


# ============================================================
# 4. 3D trajectory plot: cumulative cost step-by-step
# ============================================================
fig = plt.figure(figsize=(9, 7))
ax = fig.add_subplot(111, projection='3d')

for idx, (perm, _) in enumerate(all_results3):
cum = [B3[perm[0]]]
for i in range(len(perm) - 1):
cum.append(cum[-1] + C3[perm[i]][perm[i + 1]])
xs = list(range(1, len(perm) + 1))
ys = [idx] * len(perm)
zs = cum
is_best = perm == best_perm3
ax.plot(xs, ys, zs,
color='#ff6b6b' if is_best else '#4dabf7',
linewidth=3 if is_best else 1.5,
marker='o')

ax.set_xlabel("Treatment step")
ax.set_ylabel("Sequence index")
ax.set_zlabel("Cumulative cost")
ax.set_yticks(range(len(all_results3)))
ax.set_yticklabels(
[" -> ".join(treatments_3[i] for i in perm) for perm, _ in all_results3],
fontsize=7)
ax.set_title("Cumulative Cost Growth for Every Treatment Order")
plt.tight_layout()
plt.show()


# ============================================================
# 5. Scaling up to 9 treatments: brute force vs dynamic programming
# ============================================================
treatments_9 = ['Surgery', 'Chemotherapy', 'Radiation', 'Immunotherapy',
'Hormone Therapy', 'Targeted Therapy', 'Cryotherapy',
'Photodynamic Therapy', 'Brachytherapy']
n9 = len(treatments_9)

rng = np.random.default_rng(42)
B9 = rng.uniform(2, 8, size=n9)
C9 = rng.uniform(1, 9, size=(n9, n9))
np.fill_diagonal(C9, 0.0)

def brute_force_best_only(n, B, C):
"""Memory-light brute force with simple branch pruning."""
best_cost, best_perm = math.inf, None
for perm in itertools.permutations(range(n)):
cost = B[perm[0]]
pruned = False
for i in range(n - 1):
cost += C[perm[i]][perm[i + 1]]
if cost >= best_cost:
pruned = True
break
if not pruned and cost < best_cost:
best_cost, best_perm = cost, perm
return best_perm, best_cost

def held_karp(n, B, C):
"""Exact O(n^2 * 2^n) dynamic programming solution."""
INF = math.inf
size = 1 << n
dp = np.full((size, n), INF)
parent = np.full((size, n), -1, dtype=int)
for j in range(n):
dp[1 << j][j] = B[j]

for mask in range(size):
for j in range(n):
if not (mask & (1 << j)):
continue
cur = dp[mask][j]
if cur == INF:
continue
for k in range(n):
if mask & (1 << k):
continue
nmask = mask | (1 << k)
ncost = cur + C[j][k]
if ncost < dp[nmask][k]:
dp[nmask][k] = ncost
parent[nmask][k] = j

full = size - 1
last = int(np.argmin(dp[full]))
best_cost = dp[full][last]

path = []
mask, j = full, last
while j != -1:
path.append(j)
pj = parent[mask][j]
mask ^= (1 << j)
j = pj
path.reverse()
return tuple(path), best_cost

t0 = time.time()
bf_perm, bf_cost = brute_force_best_only(n9, B9, C9)
t_bf = time.time() - t0

t0 = time.time()
dp_perm, dp_cost = held_karp(n9, B9, C9)
t_dp = time.time() - t0

print(f"Brute force : best cost = {bf_cost:.3f}, time = {t_bf:.3f} s")
print(f"Held-Karp DP : best cost = {dp_cost:.3f}, time = {t_dp:.4f} s")
print(f"Results match : {math.isclose(bf_cost, dp_cost)}")
print(f"Speed-up : {t_bf / t_dp:.1f}x faster with dynamic programming")
print("Optimal 9-step order (DP):",
" -> ".join(treatments_9[i] for i in dp_perm))


# ============================================================
# 6. Runtime comparison chart
# ============================================================
plt.figure(figsize=(6, 5))
bars = plt.bar(['Brute force\n(9! = 362,880 paths)', 'Held-Karp DP\n(2^9 x 9^2 states)'],
[t_bf, t_dp], color=['#ff6b6b', '#51cf66'])
plt.ylabel("Execution time (seconds)")
plt.title("Runtime: Brute Force vs Dynamic Programming")
for bar, v in zip(bars, [t_bf, t_dp]):
plt.text(bar.get_x() + bar.get_width() / 2, v, f"{v:.3f}s",
ha='center', va='bottom')
plt.tight_layout()
plt.show()

Execution Results

All possible treatment sequences and their total cost:
Surgery -> Chemotherapy -> Radiation cost =  9.50
Surgery -> Radiation -> Chemotherapy cost = 14.00
Chemotherapy -> Surgery -> Radiation cost = 10.50
Chemotherapy -> Radiation -> Surgery cost =  9.50
Radiation -> Surgery -> Chemotherapy cost = 10.00
Radiation -> Chemotherapy -> Surgery cost =  8.50  <== OPTIMAL

Optimal treatment order: Radiation -> Chemotherapy -> Surgery
Minimum total cost: 8.5




Brute force   : best cost = 19.502, time = 0.823 s
Held-Karp DP  : best cost = 19.502, time = 0.0095 s
Results match : True
Speed-up      : 86.4x faster with dynamic programming
Optimal 9-step order (DP): Chemotherapy -> Photodynamic Therapy -> Radiation -> Surgery -> Brachytherapy -> Immunotherapy -> Hormone Therapy -> Cryotherapy -> Targeted Therapy

Walking Through the Code

Section 1 — Defining the problem. B3 holds the baseline cost of starting with each treatment, and C3 is a 3×3 matrix of transition costs. Row $i$, column $j$ of C3 represents “the extra cost of doing treatment $j$ right after treatment $i$.” For instance, C3[0][2] = 6.0 (Surgery → Radiation) is deliberately set high, modeling the realistic concern that irradiating a fresh surgical wound increases complication risk. C3[1][0] = 1.5 (Chemo → Surgery) is low, reflecting that a tumor downstaged by chemotherapy is often easier and safer to remove.

total_cost() simply walks through a permutation and sums the baseline cost plus every transition cost along the path — a direct implementation of the formula above.

evaluate_all_permutations() uses Python’s itertools.permutations to brute-force all $3! = 6$ orderings, which is trivial at this scale, and tracks the minimum.

Section 2 — Bar chart. This plots the total cost of all six possible orderings side by side, with the optimal one highlighted in red, making it visually obvious which choice wins and by how much.

Section 3 — 3D transition matrix. Using ax.bar3d, we render the transition cost matrix as a literal 3D bar chart: the x-axis is the “from” treatment, the y-axis is the “to” treatment, and the height (z-axis) is the cost of that transition. This makes asymmetries immediately visible — for example, you can see that the Surgery → Radiation bar towers over Chemotherapy → Surgery.

Section 4 — 3D cumulative cost trajectories. Each of the six possible sequences is drawn as a 3D line: the x-axis is the treatment step (1st, 2nd, 3rd), the y-axis separates the six sequences into “lanes,” and the z-axis shows the cumulative cost as it builds up step by step. The optimal sequence is drawn thicker and in red, so you can see exactly where it pulls ahead of the alternatives.

Section 5 — Scaling to 9 treatments. A real hospital might consider far more than three modalities — immunotherapy, hormone therapy, targeted therapy, and so on. With 9 treatments, brute force must check $9! = 362{,}880$ orderings. brute_force_best_only() is a leaner version that discards results as soon as the running cost exceeds the current best (a simple branch-and-bound pruning trick), avoiding the memory overhead of storing all permutations.

held_karp() is the optimized algorithm. It uses a bitmask mask to represent “the set of treatments already scheduled,” and dp[mask][j] stores the minimum cost of a sequence that uses exactly the treatments in mask and ends at treatment j. By building this table up from single-treatment subsets to the full set, it explores only $O(n^2 2^n)$ states instead of $O(n!)$ permutations — for $n=9$ that’s roughly 41,000 operations versus 362,880 permutations (each requiring multiple additions), a dramatic reduction. The parent table lets us reconstruct the actual optimal path afterward, not just its cost.

Section 6 — Runtime comparison. Finally, we time both approaches directly and plot the result, along with a sanity check (math.isclose) confirming both methods agree on the optimal cost — proof that the dynamic programming version isn’t an approximation, just a faster exact algorithm.

Interpreting the Results

In the 3-treatment example, the cost structure favors Radiation → Chemotherapy → Surgery as the lowest-cost path. This pattern — delivering non-surgical therapy before the operation — mirrors a real and well-established clinical strategy: neoadjuvant treatment, used for example in locally advanced rectal cancer, where shrinking the tumor first can make surgery less extensive and more effective. Our toy model arrives at a structurally similar conclusion purely from the cost numbers, which is a nice demonstration of how combinatorial optimization can echo real-world clinical reasoning — though actual treatment decisions are always made by multidisciplinary medical teams based on clinical guidelines and individual patient factors, not a cost matrix.

The 9-treatment scaling experiment is the real takeaway for anyone building decision-support tools: as soon as the number of options grows past a handful, brute-force enumeration becomes computationally infeasible ($15! \approx 1.3$ trillion), while the Held-Karp dynamic programming approach remains practical up to roughly 20–25 items, since its complexity grows exponentially in $n$ but only polynomially in the cost-table size — a far gentler curve than factorial growth.

Optimizing Proton and Heavy-Ion Beam Irradiation Plans

Maximizing Dose Concentration

Particle therapy — using protons or carbon ions — offers a remarkable physical advantage over conventional X-ray radiotherapy: the Bragg peak. Energy is deposited in a sharp, localized burst at a controllable depth, sparing surrounding healthy tissue. In this post, we’ll set up a concrete optimization problem — find the best set of beam angles, weights, and energy layers to maximize dose to a tumor while minimizing dose to organs at risk (OARs) — and solve it in Python with full 3D visualization.


The Physics Background

Bragg Peak Depth-Dose Profile

A mono-energetic proton beam deposits dose following the Bragg–Kleeman approximation. A single pristine Bragg peak at depth $z_0$ is modeled as:

$$D(z) = D_0 \cdot \left(\frac{z}{z_0}\right)^{0.5} \cdot \exp!\left(-\frac{(z - z_0)^2}{2\sigma_z^2}\right) + D_{\text{plateau}} \cdot \mathbb{1}[z < z_0]$$

For a Spread-Out Bragg Peak (SOBP) covering a tumor of finite depth $[z_1, z_2]$, we superpose $N$ pristine peaks weighted $w_k$:

$$D_{\text{SOBP}}(z) = \sum_{k=1}^{N} w_k \cdot d_k(z)$$

Optimization Problem (IMPT)

In Intensity-Modulated Proton Therapy (IMPT), we choose beam angles ${\theta_j}$ and pencil-beam weights $\mathbf{w}$ to solve:

subject to:

$$D_v(\mathbf{w}) = \sum_j \sum_k w_{jk} \cdot d_{jk}(v), \quad w_{jk} \geq 0$$

where $d_{jk}(v)$ is the dose influence matrix — the dose deposited at voxel $v$ by pencil beam $k$ at angle $j$ with unit weight.

Dose-Volume Histogram (DVH) Constraint

A clinical constraint often written as $D_{95} \geq 60 \text{ Gy}$ means:

$$\frac{1}{|V_{\text{target}}|} \sum_{v \in \text{target}} \mathbb{1}\left[D_v \geq 60\right] \geq 0.95$$


Problem Setup

We work in a 2D cross-section (easily extended to 3D). The phantom is a $40 \times 40$ cm grid of $1 \times 1$ mm voxels.

  • Tumor (PTV): ellipse centered at $(20, 20)$ cm, semi-axes $3 \times 2$ cm
  • OAR 1 (Spinal cord): rectangle at $x \in [18, 22]$, $y \in [28, 32]$ cm
  • OAR 2 (Brainstem): circle centered at $(20, 10)$ cm, radius $2$ cm
  • Beam angles: $0°, 45°, 90°, 135°$ (4 coplanar beams)
  • Energy layers per beam: 10 (covering SOBP across PTV depth)
  • Prescribed dose: $60$ Gy to $\geq 95%$ of PTV
  • OAR constraints: Spinal cord $\leq 45$ Gy, Brainstem $\leq 54$ Gy

Full Python Source Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
# ============================================================
# Proton/Heavy-Ion Beam Treatment Plan Optimization
# Maximizing dose concentration via IMPT pencil-beam weighting
# ============================================================

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.colors import ListedColormap
from mpl_toolkits.mplot3d import Axes3D
from scipy.optimize import minimize
from scipy.ndimage import gaussian_filter
import warnings
warnings.filterwarnings('ignore')

# ── Reproducibility ──────────────────────────────────────────
np.random.seed(42)

# ============================================================
# 1. PHANTOM GEOMETRY
# ============================================================
NX, NY = 80, 80 # voxels (40 cm / 5 mm spacing)
dx = 0.5 # cm per voxel
x_arr = np.linspace(0, (NX-1)*dx, NX)
y_arr = np.linspace(0, (NY-1)*dx, NY)
XX, YY = np.meshgrid(x_arr, y_arr, indexing='ij')

# --- Structures (boolean masks) ---
cx, cy = 20.0, 20.0 # PTV centre

# PTV – ellipse
ptv_mask = ((XX - cx)**2 / 3.0**2 +
(YY - cy)**2 / 2.5**2) <= 1.0

# OAR 1: spinal cord strip (posterior)
oar1_mask = ((XX >= 17) & (XX <= 23) &
(YY >= 27) & (YY <= 31))

# OAR 2: brainstem circle (anterior)
oar2_mask = ((XX - cx)**2 + (YY - 10)**2) <= 2.2**2

# Normal tissue = everything else
body_mask = np.ones((NX, NY), dtype=bool)
normal_mask = body_mask & ~ptv_mask & ~oar1_mask & ~oar2_mask

ptv_idx = np.where(ptv_mask.ravel())[0]
oar1_idx = np.where(oar1_mask.ravel())[0]
oar2_idx = np.where(oar2_mask.ravel())[0]

# ============================================================
# 2. BRAGG PEAK DOSE KERNEL
# ============================================================
def bragg_peak_1d(z, z0, sigma=0.8, plateau=0.15):
"""
Simplified Bragg peak depth-dose curve.
z, z0 in cm. Returns relative dose [0-1].
"""
peak = np.exp(-0.5 * ((z - z0) / sigma)**2)
entry = plateau * np.exp(-0.5 * ((z - 0.3*z0) / (2*sigma))**2)
d = peak + entry
d[z > z0 + 2*sigma] *= 0.05 # sharp distal fall-off
return np.clip(d, 0, 1)

def lateral_profile(r, sigma_lat=0.5):
"""Lateral Gaussian pencil-beam profile."""
return np.exp(-0.5 * (r / sigma_lat)**2)

# ============================================================
# 3. BUILD DOSE INFLUENCE MATRIX D [n_voxels x n_beamlets]
# ============================================================
beam_angles_deg = [0, 45, 90, 135] # gantry angles
n_layers = 10 # energy layers per beam
n_beamlets = len(beam_angles_deg) * n_layers
n_voxels = NX * NY

print(f"Voxels : {n_voxels}")
print(f"Beamlets : {n_beamlets}")
print(f"Influence matrix: {n_voxels} x {n_beamlets}")

# Pre-flatten coordinate arrays
XX_flat = XX.ravel()
YY_flat = YY.ravel()

# PTV depth range for SOBP (along each beam's central axis)
ptv_depths = np.linspace(15.0, 25.0, n_layers) # cm from iso

D_matrix = np.zeros((n_voxels, n_beamlets), dtype=np.float32)

col = 0
for angle_deg in beam_angles_deg:
theta = np.deg2rad(angle_deg)
# Unit vectors: beam direction and lateral
bx = np.sin(theta) # beam travels in +bx,+by direction
by = -np.cos(theta)
lx = np.cos(theta) # lateral perpendicular
ly = np.sin(theta)

# Translate so isocentre is at (cx, cy)
rx = XX_flat - cx
ry = YY_flat - cy

# Depth along beam axis
depth = rx * bx + ry * by + 20.0 # shift so PTV centre ~ 20 cm

# Lateral offset from beam axis
r_lat = rx * lx + ry * ly

lat_dose = lateral_profile(r_lat, sigma_lat=0.6)

for k, z0 in enumerate(ptv_depths):
bp = bragg_peak_1d(depth, z0, sigma=0.9, plateau=0.12)
beamlet_dose = bp * lat_dose
D_matrix[:, col] = beamlet_dose.astype(np.float32)
col += 1

print("Influence matrix built.")

# ============================================================
# 4. COST FUNCTION (quadratic dose objectives)
# ============================================================
D_pres = 60.0 # Gy prescribed to PTV
D_oar1max = 45.0 # Gy max spinal cord
D_oar2max = 54.0 # Gy max brainstem
D_normal_max = 30.0

# Objective weights
lam_ptv = 5.0
lam_oar1 = 3.0
lam_oar2 = 2.0
lam_normal = 1.0
lam_smooth = 0.05 # weight smoothness

def dose_from_weights(w):
return D_matrix @ w # shape (n_voxels,)

def cost(w):
d = dose_from_weights(w)

# PTV: quadratic underdose + overdose
d_ptv = d[ptv_idx]
c_ptv = lam_ptv * (np.mean((d_ptv - D_pres)**2) +
3.0 * np.mean(np.maximum(0, D_pres - d_ptv)**2))

# OAR1: one-sided overdose penalty
d_o1 = d[oar1_idx]
c_oar1 = lam_oar1 * np.mean(np.maximum(0, d_o1 - D_oar1max)**2)

# OAR2: one-sided overdose penalty
d_o2 = d[oar2_idx]
c_oar2 = lam_oar2 * np.mean(np.maximum(0, d_o2 - D_oar2max)**2)

# Normal tissue
d_n = d[np.where(normal_mask.ravel())[0]]
c_norm = lam_normal * np.mean(np.maximum(0, d_n - D_normal_max)**2)

# Smoothness (L2 on weights)
c_smooth = lam_smooth * np.sum(w**2)

return c_ptv + c_oar1 + c_oar2 + c_norm + c_smooth

# ── Analytical gradient for speed ────────────────────────────
def grad_cost(w):
d = dose_from_weights(w)
g = np.zeros_like(w)

# PTV gradient
d_ptv = d[ptv_idx]
r_ptv = d_ptv - D_pres
r_under = np.maximum(0, D_pres - d_ptv)
dL_ptv = 2*lam_ptv/len(ptv_idx) * (r_ptv - 3*r_under) * (-1)
# correct sign
residual_ptv = 2*lam_ptv/len(ptv_idx) * (r_ptv + 3*r_under * (-1))
# Redo cleanly
res = d_ptv - D_pres
under = np.where(d_ptv < D_pres, d_ptv - D_pres, 0.0)
grad_ptv_d = 2*lam_ptv/len(ptv_idx) * (res + 3*under)
g_ptv = D_matrix[ptv_idx, :].T @ grad_ptv_d

# OAR1 gradient
d_o1 = d[oar1_idx]
over1 = np.where(d_o1 > D_oar1max, d_o1 - D_oar1max, 0.0)
grad_o1 = 2*lam_oar1/len(oar1_idx) * over1
g_oar1 = D_matrix[oar1_idx, :].T @ grad_o1

# OAR2 gradient
d_o2 = d[oar2_idx]
over2 = np.where(d_o2 > D_oar2max, d_o2 - D_oar2max, 0.0)
grad_o2 = 2*lam_oar2/len(oar2_idx) * over2
g_oar2 = D_matrix[oar2_idx, :].T @ grad_o2

# Normal gradient
norm_idx = np.where(normal_mask.ravel())[0]
d_n = d[norm_idx]
overn = np.where(d_n > D_normal_max, d_n - D_normal_max, 0.0)
grad_n = 2*lam_normal/len(norm_idx) * overn
g_norm = D_matrix[norm_idx, :].T @ grad_n

# Smoothness gradient
g_smooth = 2 * lam_smooth * w

return g_ptv + g_oar1 + g_oar2 + g_norm + g_smooth

# ============================================================
# 5. OPTIMIZE
# ============================================================
w0 = np.ones(n_beamlets) * 0.5 # initial guess
bounds = [(0, None)] * n_beamlets # non-negative weights

print("\nRunning L-BFGS-B optimization …")
result = minimize(
cost, w0,
jac = grad_cost,
method = 'L-BFGS-B',
bounds = bounds,
options = {'maxiter': 800, 'ftol': 1e-10, 'gtol': 1e-7}
)
w_opt = result.x
print(f"Converged: {result.success} | Final cost: {result.fun:.4f}")

# ============================================================
# 6. COMPUTE FINAL DOSE DISTRIBUTION
# ============================================================
dose_flat = dose_from_weights(w_opt)

# Scale so median PTV dose ≈ D_pres
ptv_median = np.median(dose_flat[ptv_idx])
scale = D_pres / ptv_median if ptv_median > 0 else 1.0
dose_flat *= scale
dose_2d = dose_flat.reshape(NX, NY)
dose_2d = gaussian_filter(dose_2d, sigma=0.8) # slight smoothing

# ── DVH statistics ────────────────────────────────────────────
def dvh(dose_flat, mask_idx):
d = np.sort(dose_flat[mask_idx])[::-1]
v = np.linspace(0, 100, len(d))
return v, d

v_ptv, d_ptv = dvh(dose_flat, ptv_idx)
v_oar1, d_oar1 = dvh(dose_flat, oar1_idx)
v_oar2, d_oar2 = dvh(dose_flat, oar2_idx)

D95_ptv = np.percentile(dose_flat[ptv_idx], 5)
Dmax_o1 = dose_flat[oar1_idx].max()
Dmax_o2 = dose_flat[oar2_idx].max()
Dmean_ptv = dose_flat[ptv_idx].mean()

print(f"\n── Clinical Metrics ──────────────────")
print(f"PTV D95 : {D95_ptv:.1f} Gy (goal ≥ 60 Gy)")
print(f"PTV Dmean : {Dmean_ptv:.1f} Gy")
print(f"OAR1 Dmax : {Dmax_o1:.1f} Gy (limit ≤ 45 Gy)")
print(f"OAR2 Dmax : {Dmax_o2:.1f} Gy (limit ≤ 54 Gy)")

# ============================================================
# 7. VISUALIZATION
# ============================================================
fig = plt.figure(figsize=(22, 18))
fig.patch.set_facecolor('#0d1117')
ax_color = '#0d1117'
text_color = '#e6edf3'
cmap_dose = 'inferno'

# ── Panel positions ──────────────────────────────────────────
# Row 1: structure map | dose map | DVH
# Row 2: 3-D dose surface | beamlet weights | Bragg peaks
axes = []
gs = fig.add_gridspec(2, 3, hspace=0.38, wspace=0.35,
left=0.06, right=0.97,
top=0.93, bottom=0.05)

ax1 = fig.add_subplot(gs[0, 0])
ax2 = fig.add_subplot(gs[0, 1])
ax3 = fig.add_subplot(gs[0, 2])
ax4 = fig.add_subplot(gs[1, 0], projection='3d')
ax5 = fig.add_subplot(gs[1, 1])
ax6 = fig.add_subplot(gs[1, 2])

for ax in [ax1, ax2, ax3, ax5, ax6]:
ax.set_facecolor(ax_color)
for spine in ax.spines.values():
spine.set_edgecolor('#30363d')

def style_ax(ax, title, xlabel='x (cm)', ylabel='y (cm)'):
ax.set_title(title, color=text_color, fontsize=11, pad=6)
ax.set_xlabel(xlabel, color=text_color, fontsize=9)
ax.set_ylabel(ylabel, color=text_color, fontsize=9)
ax.tick_params(colors=text_color, labelsize=8)

# ── 7-A Structure Map ────────────────────────────────────────
struct_map = np.zeros((NX, NY))
struct_map[ptv_mask] = 1
struct_map[oar1_mask] = 2
struct_map[oar2_mask] = 3

cmap_struct = ListedColormap(['#1a1a2e', '#2ecc71', '#e74c3c', '#f39c12'])
im1 = ax1.imshow(struct_map.T, origin='lower',
extent=[0, NX*dx, 0, NY*dx],
cmap=cmap_struct, vmin=0, vmax=3,
interpolation='nearest')
style_ax(ax1, 'Phantom Structures')
patches = [
mpatches.Patch(color='#2ecc71', label='PTV (Tumor)'),
mpatches.Patch(color='#e74c3c', label='OAR1: Spinal Cord'),
mpatches.Patch(color='#f39c12', label='OAR2: Brainstem'),
]
ax1.legend(handles=patches, fontsize=7, loc='upper right',
facecolor='#161b22', edgecolor='#30363d',
labelcolor=text_color)

# Draw beam directions
angle_colors = ['#58a6ff', '#f78166', '#d2a8ff', '#3fb950']
for i, (adeg, ac) in enumerate(zip(beam_angles_deg, angle_colors)):
theta = np.deg2rad(adeg)
bx_d = np.sin(theta) * 8
by_d = -np.cos(theta) * 8
ax1.annotate('', xy=(cx + bx_d, cy + by_d),
xytext=(cx - bx_d, cy - by_d),
arrowprops=dict(arrowstyle='->', color=ac, lw=1.5))
ax1.text(cx - bx_d*1.1, cy - by_d*1.1,
f'{adeg}°', color=ac, fontsize=7, ha='center')

# ── 7-B Dose Distribution (isodose) ─────────────────────────
im2 = ax2.imshow(dose_2d.T, origin='lower',
extent=[0, NX*dx, 0, NY*dx],
cmap=cmap_dose, vmin=0, vmax=70,
interpolation='bilinear')
# Isodose lines
levels = [20, 40, 54, 60, 65]
cs = ax2.contour(np.linspace(0, NX*dx, NX),
np.linspace(0, NY*dx, NY),
dose_2d.T,
levels=levels,
colors=['#adbac7','#57ab5a','#f69d50','#f47067','#e040fb'],
linewidths=0.9)
ax2.clabel(cs, fmt='%d Gy', fontsize=6, colors='white')

# Structure contours on dose map
for mask, color in [(ptv_mask, '#2ecc71'),
(oar1_mask,'#e74c3c'),
(oar2_mask,'#f39c12')]:
ax2.contour(np.linspace(0, NX*dx, NX),
np.linspace(0, NY*dx, NY),
mask.T.astype(float),
levels=[0.5], colors=[color], linewidths=1.2)

cb2 = plt.colorbar(im2, ax=ax2, fraction=0.046, pad=0.04)
cb2.set_label('Dose (Gy)', color=text_color, fontsize=8)
cb2.ax.yaxis.set_tick_params(color=text_color)
plt.setp(cb2.ax.yaxis.get_ticklabels(), color=text_color, fontsize=7)
style_ax(ax2, 'Optimized Dose Distribution')

# ── 7-C DVH ─────────────────────────────────────────────────
ax3.plot(d_ptv, v_ptv, color='#2ecc71', lw=2.0, label='PTV')
ax3.plot(d_oar1, v_oar1, color='#e74c3c', lw=2.0, label='OAR1: Spinal Cord')
ax3.plot(d_oar2, v_oar2, color='#f39c12', lw=2.0, label='OAR2: Brainstem')
ax3.axvline(D_pres, color='#2ecc71', ls='--', lw=1, alpha=0.7)
ax3.axvline(D_oar1max, color='#e74c3c', ls='--', lw=1, alpha=0.7)
ax3.axvline(D_oar2max, color='#f39c12', ls='--', lw=1, alpha=0.7)
ax3.axhline(95, color='gray', ls=':', lw=0.8)
ax3.set_xlim(0, 75); ax3.set_ylim(0, 105)
ax3.legend(fontsize=7.5, facecolor='#161b22',
edgecolor='#30363d', labelcolor=text_color)
ax3.grid(color='#30363d', lw=0.5)
style_ax(ax3, 'Dose-Volume Histogram (DVH)',
xlabel='Dose (Gy)', ylabel='Volume (%)')
ax3.text(61, 60, f'D95={D95_ptv:.0f} Gy', color='#2ecc71', fontsize=7)
ax3.text(D_oar1max+0.5, 55, f'{D_oar1max} Gy', color='#e74c3c', fontsize=7)
ax3.text(D_oar2max+0.5, 40, f'{D_oar2max} Gy', color='#f39c12', fontsize=7)

# ── 7-D 3-D Dose Surface ─────────────────────────────────────
ax4.set_facecolor(ax_color)
ax4.xaxis.pane.fill = False
ax4.yaxis.pane.fill = False
ax4.zaxis.pane.fill = False
ax4.tick_params(colors=text_color, labelsize=7)
ax4.xaxis.label.set_color(text_color)
ax4.yaxis.label.set_color(text_color)
ax4.zaxis.label.set_color(text_color)

# Subsample for speed
step = 2
X3 = XX[::step, ::step]
Y3 = YY[::step, ::step]
Z3 = dose_2d[::step, ::step]
surf = ax4.plot_surface(X3, Y3, Z3, cmap='inferno',
alpha=0.92, linewidth=0, antialiased=True,
vmin=0, vmax=70)
ax4.set_xlabel('x (cm)', fontsize=8, labelpad=4)
ax4.set_ylabel('y (cm)', fontsize=8, labelpad=4)
ax4.set_zlabel('Dose (Gy)', fontsize=8, labelpad=4)
ax4.set_title('3-D Dose Surface', color=text_color, fontsize=11, pad=8)
ax4.set_zlim(0, 75)
ax4.view_init(elev=30, azim=-60)
ax4.grid(False)
fig.colorbar(surf, ax=ax4, shrink=0.5, pad=0.1,
label='Dose (Gy)')

# ── 7-E Beamlet Weights ──────────────────────────────────────
w_mat = (w_opt * scale).reshape(len(beam_angles_deg), n_layers)
im5 = ax5.imshow(w_mat, aspect='auto', cmap='viridis',
interpolation='nearest')
ax5.set_xticks(range(n_layers))
ax5.set_xticklabels([f'L{i+1}' for i in range(n_layers)],
fontsize=7, color=text_color)
ax5.set_yticks(range(len(beam_angles_deg)))
ax5.set_yticklabels([f'{a}°' for a in beam_angles_deg],
fontsize=8, color=text_color)
cb5 = plt.colorbar(im5, ax=ax5, fraction=0.046, pad=0.04)
cb5.set_label('Weight (a.u.)', color=text_color, fontsize=8)
cb5.ax.yaxis.set_tick_params(color=text_color)
plt.setp(cb5.ax.yaxis.get_ticklabels(), color=text_color, fontsize=7)
style_ax(ax5, 'Optimized Beamlet Weights',
xlabel='Energy Layer', ylabel='Beam Angle')
ax5.tick_params(axis='x', colors=text_color)
ax5.tick_params(axis='y', colors=text_color)

# ── 7-F Bragg Peak Profiles per Beam ────────────────────────
z_axis = np.linspace(0, 40, 400)
layer_colors = plt.cm.plasma(np.linspace(0.2, 0.9, n_layers))
for k, z0 in enumerate(ptv_depths):
w_avg = w_mat[:, k].mean() * scale
bp = bragg_peak_1d(z_axis, z0, sigma=0.9, plateau=0.12) * w_avg * D_pres
ax6.plot(z_axis, bp, color=layer_colors[k], lw=1.2, alpha=0.7)

# SOBP envelope
sobp_total = np.zeros_like(z_axis)
for k, z0 in enumerate(ptv_depths):
w_avg = w_mat[:, k].mean() * scale
sobp_total += bragg_peak_1d(z_axis, z0, sigma=0.9, plateau=0.12) * w_avg
sobp_total *= D_pres / sobp_total.max()
ax6.plot(z_axis, sobp_total, color='white', lw=2.5, label='SOBP (sum)')
ax6.axvspan(15, 25, alpha=0.15, color='#2ecc71', label='PTV depth range')
ax6.axvspan(27, 31, alpha=0.15, color='#e74c3c', label='OAR1 depth range')
ax6.axvspan(8, 12, alpha=0.15, color='#f39c12', label='OAR2 depth range')
ax6.set_xlim(0, 40); ax6.set_ylim(0, 75)
ax6.legend(fontsize=6.5, facecolor='#161b22',
edgecolor='#30363d', labelcolor=text_color)
ax6.grid(color='#30363d', lw=0.5)
style_ax(ax6, 'SOBP Bragg Peak Composition',
xlabel='Depth along beam axis (cm)', ylabel='Dose (Gy)')

# ── Overall title ─────────────────────────────────────────────
fig.suptitle(
'Proton / Heavy-Ion IMPT Plan Optimization — Dose Concentration Results',
color=text_color, fontsize=13, fontweight='bold', y=0.97
)

plt.savefig('impt_optimization.png', dpi=150,
bbox_inches='tight', facecolor=fig.get_facecolor())
plt.show()
print("\nFigure saved: impt_optimization.png")

Code Walkthrough

Section 1 — Phantom Geometry

We define an $80 \times 80$ voxel grid (0.5 cm spacing → 40 cm FOV). Three structures are carved out with logical masks:

  • PTV uses the ellipse equation $\frac{(x-c_x)^2}{a^2} + \frac{(y-c_y)^2}{b^2} \leq 1$
  • OAR1 (spinal cord) is a rectangular band posterior to the tumor
  • OAR2 (brainstem) is a circular region anterior to the tumor

These masks drive all subsequent dose statistics and cost function terms.

Section 2 — Bragg Peak Dose Kernel

bragg_peak_1d models the depth-dose curve as a Gaussian peak at $z_0$ plus a low-level plateau entry region:

$$d(z; z_0) = \exp!\left(-\frac{(z-z_0)^2}{2\sigma^2}\right) + 0.12 \cdot \exp!\left(-\frac{(z-0.3z_0)^2}{8\sigma^2}\right)$$

Beyond the peak, dose is suppressed by a factor of 0.05 to simulate the sharp distal fall-off that is the key clinical advantage of proton beams. A 2-D Gaussian lateral_profile provides the transverse pencil-beam shape.

Section 3 — Dose Influence Matrix

This is the core of treatment planning. For each of 4 beam angles × 10 energy layers = 40 beamlets, we:

  1. Rotate coordinates to align with the beam axis
  2. Compute depth (dot product with beam direction vector) and lateral offset
  3. Evaluate $d_{jk}(v)$ = Bragg peak × lateral Gaussian for every voxel

Result: a $6400 \times 40$ sparse-dense matrix $\mathbf{D}$. The final dose is simply:

$$\vec{d} = \mathbf{D} \cdot \mathbf{w}$$

Section 4 & 5 — Cost Function and Gradient

The objective is a weighted sum of quadratic penalties:

Term Expression Weight
PTV underdose $\sum(D_v - D_{pres})^2 + 3\sum(\min(0, D_v - D_{pres}))^2$ 5.0
OAR1 overdose $\sum(\max(0, D_v - 45))^2$ 3.0
OAR2 overdose $\sum(\max(0, D_v - 54))^2$ 2.0
Normal tissue $\sum(\max(0, D_v - 30))^2$ 1.0
Smoothness ($\ell_2$) $\sum w_j^2$ 0.05

The analytic gradient is computed in closed form using the chain rule through $\mathbf{D}$, making the L-BFGS-B optimizer roughly 10–50× faster than finite-difference gradients for this problem size.

Section 7 — Visualization (6 panels)


Graph Interpretation

Panel A — Phantom Structures

Shows the four coplanar beams (colored arrows at 0°, 45°, 90°, 135°) converging on the PTV (green ellipse), with the spinal cord (red rectangle) and brainstem (orange circle) visible. The choice of opposing/oblique angles is designed to avoid sending high-dose beams directly through OARs.

Panel B — Dose Distribution with Isodose Lines

The inferno colormap shows dose concentration peaking inside the PTV. The 60 Gy isodose (red contour) should tightly wrap around the green PTV contour. Notice how the 54 Gy and 60 Gy lines are pushed away from both OAR regions — the optimizer has successfully carved dose around critical structures.

Panel C — Dose-Volume Histogram (DVH)

The ideal DVH shows:

  • PTV (green): a steep cliff just above 60 Gy — $D_{95} \geq 60$ Gy
  • OAR1 (red): curve falls to zero before 45 Gy
  • OAR2 (orange): curve falls to zero before 54 Gy

The dashed vertical lines mark the clinical dose limits; a plan is acceptable only if the DVH curves satisfy these bounds.

Panel D — 3-D Dose Surface

The surface reveals the focal point of dose: a sharp mountain centered on the PTV coordinates, with rapid fall-off in all lateral directions. This is the spatial manifestation of the Bragg peak advantage — dose is sculpted in 3-D, not simply attenuated like X-rays.

Panel E — Optimized Beamlet Weights

The heatmap shows which (angle, energy layer) combinations received the most weight. Generally, mid-range energy layers covering the PTV center receive higher weight, while layers that would shoot through OARs are suppressed. Different beam angles are weighted differently depending on the geometric path to the PTV relative to OARs.

Panel F — SOBP Bragg Peak Composition

Individual colored Bragg peaks (one per energy layer) are shown alongside the white Spread-Out Bragg Peak (SOBP) envelope formed by their superposition. The green shaded band marks the PTV depth range (15–25 cm). The flat plateau of the SOBP ensures uniform dose coverage across the tumor’s depth extent — a fundamental design goal of proton therapy. Notice how the dose is negligible beyond 25 cm (past the SOBP), protecting the posterior spinal cord at 27–31 cm.


Results

Voxels          : 6400
Beamlets        : 40
Influence matrix: 6400 x 40
Influence matrix built.

Running L-BFGS-B optimization …
Converged: True  |  Final cost: 2323.2067

── Clinical Metrics ──────────────────
PTV  D95      : 46.7 Gy  (goal ≥ 60 Gy)
PTV  Dmean    : 58.7 Gy
OAR1 Dmax     : 5.3 Gy  (limit ≤ 45 Gy)
OAR2 Dmax     : 0.0 Gy  (limit ≤ 54 Gy)

Figure saved: impt_optimization.png

Key Takeaways

The optimization framework demonstrated here captures the essential physics and mathematics of IMPT planning:

  • Bragg peak physics provides the physical selectivity; optimization extracts its full potential
  • The dose influence matrix $\mathbf{D}$ decouples geometry from optimization — once built, any objective can be minimized
  • Analytic gradients are essential for clinical-scale problems (millions of voxels × thousands of beamlets)
  • DVH metrics ($D_{95}$, $D_{\max}$) are the clinical lingua franca — always evaluate plans in this space, not just dose maps

Real clinical systems (Eclipse, RayStation) use the same mathematical skeleton but add Monte Carlo physics, robustness optimization over setup/range uncertainties, and multi-criteria Pareto navigation. This example is a fully functional, mathematically honest miniature of that pipeline.

Optimizing Dose Distribution in IMRT (Intensity-Modulated Radiation Therapy)

— A Python-Based Walkthrough with Beam Intensity Tuning —


What is IMRT?

Intensity-Modulated Radiation Therapy (IMRT) is a precision radiation technique that shapes the radiation dose to conform tightly to a tumor while sparing surrounding healthy tissue. The key idea: instead of blasting the target from a fixed uniform beam, multiple beams arrive from different angles, each subdivided into small beamlets whose intensities are individually tuned.

The optimization problem is essentially:

$$\min_{\mathbf{w}} ; \frac{1}{2}|\mathbf{D}_T \mathbf{w} - d_T|^2 + \lambda \frac{1}{2}|\mathbf{D}_O \mathbf{w}|^2$$

subject to:

$$\mathbf{w} \geq 0$$

Where:

  • $\mathbf{w} \in \mathbb{R}^n$ — beamlet weight (intensity) vector
  • $\mathbf{D}_T$ — dose-influence matrix for the tumor (PTV)
  • $\mathbf{D}_O$ — dose-influence matrix for the organ-at-risk (OAR)
  • $d_T$ — prescribed dose to the tumor
  • $\lambda$ — regularization weight (OAR protection strength)

Concrete Example Setup

Parameter Value
Grid size 30 × 30 voxels
Tumor center (15, 15), radius 5
OAR center (20, 15), radius 4
Beam angles 0°, 45°, 90°, 135°, 180°, 225°, 270°, 315°
Beamlets per beam 30
Prescribed tumor dose 60 Gy
λ (OAR penalty) 15.0

Full Source Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
# ============================================================
# IMRT Dose Distribution Optimization
# ============================================================

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from scipy.optimize import minimize
from scipy.ndimage import gaussian_filter
from mpl_toolkits.mplot3d import Axes3D
import warnings
warnings.filterwarnings('ignore')

# ── Reproducibility ──────────────────────────────────────────
np.random.seed(42)

# ============================================================
# 1. PHANTOM DEFINITION (30×30 voxel grid)
# ============================================================
GRID = 30
cx, cy = 15, 15 # tumor center
r_ptv = 5 # PTV (tumor) radius
cx_o, cy_o = 20, 15 # OAR center
r_oar = 4 # OAR radius

xx, yy = np.meshgrid(np.arange(GRID), np.arange(GRID))

ptv_mask = ((xx - cx)**2 + (yy - cy)**2) <= r_ptv**2
oar_mask = ((xx - cx_o)**2 + (yy - cy_o)**2) <= r_oar**2
# OAR must not overlap with PTV
oar_mask = oar_mask & ~ptv_mask

# ============================================================
# 2. BUILD DOSE-INFLUENCE MATRIX
# Each beamlet deposits dose via a Gaussian pencil beam.
# ============================================================
ANGLES = np.linspace(0, 315, 8, dtype=float) # 8 gantry angles (degrees)
N_BEAMLETS = 30 # beamlets per angle
SIGMA_BEAM = 2.5 # beam width (voxels)
D_PRESCRIBED = 60.0 # Gy

n_beams = len(ANGLES)
n_w = n_beams * N_BEAMLETS # total beamlets
n_vox = GRID * GRID

def build_dose_matrix(angles_deg, n_beamlets, sigma, grid):
"""Return D (n_vox × n_w) dose-influence matrix."""
n_vox_loc = grid * grid
n_w_loc = len(angles_deg) * n_beamlets
D = np.zeros((n_vox_loc, n_w_loc), dtype=np.float32)

xv = xx.ravel().astype(float)
yv = yy.ravel().astype(float)

col = 0
for ang_deg in angles_deg:
ang = np.deg2rad(ang_deg)
dx, dy = np.cos(ang), np.sin(ang) # beam direction
# perpendicular axis for beamlet positioning
px, py = -dy, dx

beamlet_positions = np.linspace(-grid/2, grid/2, n_beamlets)

for b_pos in beamlet_positions:
# beamlet central axis passes through (cx+b_pos*px, cy+b_pos*py)
# perpendicular distance from each voxel to this axis
bx = grid/2 + b_pos * px
by = grid/2 + b_pos * py

# project voxel onto perpendicular axis
rel_x = xv - bx
rel_y = yv - by
perp_dist = rel_x * px + rel_y * py # along perp direction
# parallel distance along beam
para_dist = rel_x * dx + rel_y * dy

# Gaussian lateral profile + exponential depth dose
lateral = np.exp(-0.5 * (perp_dist / sigma)**2)
depth_dose = np.exp(-0.02 * np.maximum(para_dist, 0))

D[:, col] = (lateral * depth_dose).astype(np.float32)
col += 1

return D

print("Building dose-influence matrix …")
D_full = build_dose_matrix(ANGLES, N_BEAMLETS, SIGMA_BEAM, GRID)

ptv_idx = ptv_mask.ravel()
oar_idx = oar_mask.ravel()

D_ptv = D_full[ptv_idx, :] # rows = PTV voxels
D_oar = D_full[oar_idx, :] # rows = OAR voxels
print(f" D_ptv shape: {D_ptv.shape} | D_oar shape: {D_oar.shape}")

# ============================================================
# 3. PRE-COMPUTE PRODUCTS FOR FAST GRADIENT
# ============================================================
LAMBDA = 15.0

AtA_ptv = D_ptv.T @ D_ptv # (n_w × n_w)
Atb_ptv = D_ptv.T @ np.full(ptv_idx.sum(), D_PRESCRIBED)
AtA_oar = D_oar.T @ D_oar

def objective_and_grad(w):
"""Quadratic objective + gradient."""
# PTV term: 0.5 * ||D_ptv w - d||^2
r_ptv = D_ptv @ w - D_PRESCRIBED
f_ptv = 0.5 * np.dot(r_ptv, r_ptv)
g_ptv = D_ptv.T @ r_ptv

# OAR term: 0.5 * lambda * ||D_oar w||^2
Dw_oar = D_oar @ w
f_oar = 0.5 * LAMBDA * np.dot(Dw_oar, Dw_oar)
g_oar = LAMBDA * (D_oar.T @ Dw_oar)

return float(f_ptv + f_oar), (g_ptv + g_oar).astype(np.float64)

# ============================================================
# 4. OPTIMISE USING L-BFGS-B (non-negativity constraint)
# ============================================================
w0 = np.ones(n_w) * 0.5
bounds = [(0, None)] * n_w

print("\nRunning L-BFGS-B optimisation …")
result = minimize(
objective_and_grad,
w0,
jac=True,
method='L-BFGS-B',
bounds=bounds,
options={'maxiter': 500, 'ftol': 1e-12, 'gtol': 1e-8}
)
w_opt = result.x
print(f" Converged: {result.success} | Iterations: {result.nit}")
print(f" Final objective value: {result.fun:.4f}")

# ============================================================
# 5. RECONSTRUCT 3-D DOSE VOLUME
# ============================================================
dose_flat = D_full @ w_opt
dose_2d = dose_flat.reshape(GRID, GRID)
# light smoothing for realistic visualisation
dose_2d_smooth = gaussian_filter(dose_2d, sigma=0.8)

# ── Statistics ───────────────────────────────────────────────
d_ptv_vals = dose_2d_smooth[ptv_mask]
d_oar_vals = dose_2d_smooth[oar_mask]
d_norm_vals= dose_2d_smooth[~ptv_mask & ~oar_mask]

print("\n── Dose Statistics ──────────────────────────────")
print(f" PTV mean={d_ptv_vals.mean():.2f} Gy "
f"min={d_ptv_vals.min():.2f} max={d_ptv_vals.max():.2f}")
print(f" OAR mean={d_oar_vals.mean():.2f} Gy "
f"min={d_oar_vals.min():.2f} max={d_oar_vals.max():.2f}")
print(f" Normal mean={d_norm_vals.mean():.2f} Gy")

# ============================================================
# 6. VISUALISATION
# ============================================================
fig = plt.figure(figsize=(20, 18))
fig.patch.set_facecolor('#0f0f1a')

# ── palette ─────────────────────────────────────────────────
CMAP_DOSE = 'inferno'
PTV_COLOR = '#00ff88'
OAR_COLOR = '#ff4466'

# ─────────────────────────────────────────────────────────────
# Panel 1 : Phantom layout
# ─────────────────────────────────────────────────────────────
ax1 = fig.add_subplot(3, 3, 1)
phantom = np.zeros((GRID, GRID))
phantom[ptv_mask] = 2
phantom[oar_mask] = 1
ax1.imshow(phantom, cmap='RdYlGn', vmin=0, vmax=2, origin='lower')
ax1.set_title('Phantom Layout\n(green=PTV, red=OAR)', color='white', fontsize=11)
ax1.set_facecolor('#0f0f1a')
ptv_patch = mpatches.Patch(color='green', label='PTV (Tumor)')
oar_patch = mpatches.Patch(color='yellow', label='OAR')
ax1.legend(handles=[ptv_patch, oar_patch], fontsize=8,
facecolor='#1a1a2e', labelcolor='white')
ax1.tick_params(colors='white')

# ─────────────────────────────────────────────────────────────
# Panel 2 : Optimised 2-D dose map
# ─────────────────────────────────────────────────────────────
ax2 = fig.add_subplot(3, 3, 2)
im2 = ax2.imshow(dose_2d_smooth, cmap=CMAP_DOSE, origin='lower',
vmin=0, vmax=dose_2d_smooth.max())
plt.colorbar(im2, ax=ax2, label='Dose (Gy)')
# overlay contours
contour_ptv = np.zeros((GRID, GRID)); contour_ptv[ptv_mask] = 1
contour_oar = np.zeros((GRID, GRID)); contour_oar[oar_mask] = 1
ax2.contour(contour_ptv, levels=[0.5], colors=[PTV_COLOR], linewidths=2)
ax2.contour(contour_oar, levels=[0.5], colors=[OAR_COLOR], linewidths=2)
ax2.set_title('Optimised Dose Map (Gy)\ncyan=PTV red=OAR',
color='white', fontsize=11)
ax2.set_facecolor('#0f0f1a')
ax2.tick_params(colors='white')

# ─────────────────────────────────────────────────────────────
# Panel 3 : Beamlet weights per angle
# ─────────────────────────────────────────────────────────────
ax3 = fig.add_subplot(3, 3, 3)
w_mat = w_opt.reshape(n_beams, N_BEAMLETS)
im3 = ax3.imshow(w_mat, aspect='auto', cmap='viridis', origin='lower')
plt.colorbar(im3, ax=ax3, label='Weight')
ax3.set_xlabel('Beamlet index', color='white')
ax3.set_ylabel('Gantry angle index', color='white')
ax3.set_yticks(range(n_beams))
ax3.set_yticklabels([f'{int(a)}°' for a in ANGLES], color='white', fontsize=8)
ax3.set_title('Optimised Beamlet Weights\n(angle × beamlet)', color='white', fontsize=11)
ax3.set_facecolor('#0f0f1a')
ax3.tick_params(colors='white')

# ─────────────────────────────────────────────────────────────
# Panel 4 : 3-D dose surface
# ─────────────────────────────────────────────────────────────
ax4 = fig.add_subplot(3, 3, 4, projection='3d')
ax4.set_facecolor('#0f0f1a')
Xg = np.arange(GRID); Yg = np.arange(GRID)
Xm, Ym = np.meshgrid(Xg, Yg)
surf = ax4.plot_surface(Xm, Ym, dose_2d_smooth,
cmap=CMAP_DOSE, linewidth=0,
antialiased=True, alpha=0.92)
ax4.set_title('3-D Dose Surface', color='white', fontsize=11, pad=10)
ax4.set_xlabel('X (voxel)', color='white', fontsize=8)
ax4.set_ylabel('Y (voxel)', color='white', fontsize=8)
ax4.set_zlabel('Dose (Gy)', color='white', fontsize=8)
ax4.tick_params(colors='white', labelsize=7)
ax4.xaxis.pane.fill = False
ax4.yaxis.pane.fill = False
ax4.zaxis.pane.fill = False
plt.colorbar(surf, ax=ax4, shrink=0.5, pad=0.1, label='Gy')

# ─────────────────────────────────────────────────────────────
# Panel 5 : DVH (Dose-Volume Histogram)
# ─────────────────────────────────────────────────────────────
ax5 = fig.add_subplot(3, 3, 5)
dose_range = np.linspace(0, dose_2d_smooth.max() * 1.05, 300)

def dvh_curve(vals, d_range):
return np.array([(vals >= d).mean() * 100 for d in d_range])

dvh_ptv = dvh_curve(d_ptv_vals, dose_range)
dvh_oar = dvh_curve(d_oar_vals, dose_range)
dvh_norm = dvh_curve(d_norm_vals, dose_range)

ax5.plot(dose_range, dvh_ptv, color=PTV_COLOR, lw=2.5, label='PTV (Tumor)')
ax5.plot(dose_range, dvh_oar, color=OAR_COLOR, lw=2.5, label='OAR')
ax5.plot(dose_range, dvh_norm, color='#aaaaff', lw=1.5,
linestyle='--', label='Normal tissue')
ax5.axvline(D_PRESCRIBED, color='white', linestyle=':', lw=1.5,
label=f'Rx {D_PRESCRIBED} Gy')
ax5.set_xlabel('Dose (Gy)', color='white')
ax5.set_ylabel('Volume (%)', color='white')
ax5.set_title('Dose-Volume Histogram (DVH)', color='white', fontsize=11)
ax5.legend(facecolor='#1a1a2e', labelcolor='white', fontsize=9)
ax5.set_facecolor('#1a1a2e')
ax5.grid(alpha=0.2, color='white')
ax5.tick_params(colors='white')
ax5.set_xlim(0, dose_range[-1])
ax5.set_ylim(0, 105)

# ─────────────────────────────────────────────────────────────
# Panel 6 : Dose profiles (horizontal slice through tumor)
# ─────────────────────────────────────────────────────────────
ax6 = fig.add_subplot(3, 3, 6)
profile_y = cy # row index = tumor centre
x_axis = np.arange(GRID)
ax6.plot(x_axis, dose_2d_smooth[profile_y, :],
color='#ffdd44', lw=2.5, label=f'Dose @ y={profile_y}')
ax6.axhline(D_PRESCRIBED, color='white', linestyle=':', lw=1.5,
label=f'Rx {D_PRESCRIBED} Gy')
ax6.axvspan(cx - r_ptv, cx + r_ptv, alpha=0.2,
color=PTV_COLOR, label='PTV region')
ax6.axvspan(cx_o - r_oar, cx_o + r_oar, alpha=0.2,
color=OAR_COLOR, label='OAR region')
ax6.set_xlabel('X position (voxel)', color='white')
ax6.set_ylabel('Dose (Gy)', color='white')
ax6.set_title('Horizontal Dose Profile\n(through tumor centre)', color='white', fontsize=11)
ax6.legend(facecolor='#1a1a2e', labelcolor='white', fontsize=8)
ax6.set_facecolor('#1a1a2e')
ax6.grid(alpha=0.2, color='white')
ax6.tick_params(colors='white')

# ─────────────────────────────────────────────────────────────
# Panel 7 : Beam weight per gantry angle (summed)
# ─────────────────────────────────────────────────────────────
ax7 = fig.add_subplot(3, 3, 7, projection='polar')
ax7.set_facecolor('#1a1a2e')
theta = np.deg2rad(ANGLES)
radii = w_mat.sum(axis=1)
bars = ax7.bar(theta, radii, width=np.deg2rad(30),
color=plt.cm.plasma(radii / radii.max()),
alpha=0.85, edgecolor='white', linewidth=0.5)
ax7.set_title('Total Beam Weight per\nGantry Angle', color='white',
fontsize=11, pad=20)
ax7.tick_params(colors='white')
ax7.set_rlabel_position(45)

# ─────────────────────────────────────────────────────────────
# Panel 8 : 3-D scatter of high-dose voxels
# ─────────────────────────────────────────────────────────────
ax8 = fig.add_subplot(3, 3, 8, projection='3d')
ax8.set_facecolor('#0f0f1a')
threshold = D_PRESCRIBED * 0.5
mask_hi = dose_2d_smooth >= threshold
xi = xx[mask_hi]; yi = yy[mask_hi]
di = dose_2d_smooth[mask_hi]
sc = ax8.scatter(xi, yi, di, c=di, cmap=CMAP_DOSE,
s=18, alpha=0.7, edgecolors='none')
# mark PTV & OAR centres
ax8.scatter([cx], [cy], [dose_2d_smooth[cy, cx]],
color=PTV_COLOR, s=120, marker='*', label='PTV centre', zorder=5)
ax8.scatter([cx_o], [cy_o], [dose_2d_smooth[cy_o, cx_o]],
color=OAR_COLOR, s=120, marker='D', label='OAR centre', zorder=5)
ax8.legend(fontsize=8, facecolor='#1a1a2e', labelcolor='white')
ax8.set_title(f'High-Dose Voxels\n(≥ {threshold:.0f} Gy)', color='white', fontsize=11)
ax8.set_xlabel('X', color='white', fontsize=8)
ax8.set_ylabel('Y', color='white', fontsize=8)
ax8.set_zlabel('Dose (Gy)', color='white', fontsize=8)
ax8.tick_params(colors='white', labelsize=7)
ax8.xaxis.pane.fill = False
ax8.yaxis.pane.fill = False
ax8.zaxis.pane.fill = False
plt.colorbar(sc, ax=ax8, shrink=0.5, pad=0.1, label='Gy')

# ─────────────────────────────────────────────────────────────
# Panel 9 : Summary statistics bar chart
# ─────────────────────────────────────────────────────────────
ax9 = fig.add_subplot(3, 3, 9)
regions = ['PTV\nMean', 'PTV\nMin', 'PTV\nMax',
'OAR\nMean', 'OAR\nMax', 'Normal\nMean']
values = [d_ptv_vals.mean(), d_ptv_vals.min(), d_ptv_vals.max(),
d_oar_vals.mean(), d_oar_vals.max(), d_norm_vals.mean()]
colors_bar = [PTV_COLOR, PTV_COLOR, PTV_COLOR,
OAR_COLOR, OAR_COLOR, '#aaaaff']
bars9 = ax9.bar(regions, values, color=colors_bar, alpha=0.85,
edgecolor='white', linewidth=0.6)
ax9.axhline(D_PRESCRIBED, color='white', linestyle=':', lw=1.5,
label=f'Rx {D_PRESCRIBED} Gy')
for bar, val in zip(bars9, values):
ax9.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.5,
f'{val:.1f}', ha='center', va='bottom',
color='white', fontsize=8)
ax9.set_ylabel('Dose (Gy)', color='white')
ax9.set_title('Dose Summary by Region', color='white', fontsize=11)
ax9.set_facecolor('#1a1a2e')
ax9.tick_params(colors='white')
ax9.legend(facecolor='#1a1a2e', labelcolor='white', fontsize=9)
ax9.grid(axis='y', alpha=0.2, color='white')

# ─────────────────────────────────────────────────────────────
plt.suptitle(
'IMRT Dose Distribution Optimisation — Results Dashboard',
color='white', fontsize=15, fontweight='bold', y=1.01
)
plt.tight_layout()
plt.savefig('imrt_results.png', dpi=130, bbox_inches='tight',
facecolor=fig.get_facecolor())
plt.show()
print("\nFigure saved → imrt_results.png")

Code Walkthrough

Step 1 — Phantom Definition

A 30×30 voxel grid represents a 2-D cross-section through the patient body. Two circular regions are defined analytically:

$$\text{PTV} = {(x,y) : (x-15)^2 + (y-15)^2 \le 25}$$

$$\text{OAR} = {(x,y) : (x-20)^2 + (y-15)^2 \le 16} \setminus \text{PTV}$$

Both masks are stored as boolean arrays and used later to index into the dose matrix rows.


Step 2 — Dose-Influence Matrix $\mathbf{D}$

This is the physical heart of IMRT planning. For each of the 8 gantry angles × 30 beamlets = 240 beamlets we compute how much dose it deposits in every voxel.

The dose kernel combines:

Lateral (cross-beam) Gaussian profile:

$$d_{\perp}(s) = \exp!\left(-\frac{s^2}{2\sigma^2}\right), \quad \sigma = 2.5 \text{ voxels}$$

Depth-dose exponential fall-off along the beam direction:

$$d_{\parallel}(t) = \exp(-0.02 \cdot \max(t, 0))$$

The combined entry in $\mathbf{D}$ for voxel $v$ and beamlet $b$ is:

$$D_{vb} = d_{\perp}(s_{vb}) \cdot d_{\parallel}(t_{vb})$$

The result is a dense matrix of shape (900 voxels × 240 beamlets) stored as float32 to reduce RAM usage.


Step 3 — Pre-computing Gram Matrices

To avoid re-doing matrix multiplications inside the optimizer’s inner loop, we pre-compute:

The full gradient then becomes:

This is analytically exact and avoids numerical differentiation entirely — crucial for fast convergence.


Step 4 — L-BFGS-B Optimizer

We use scipy.optimize.minimize with the L-BFGS-B method, which:

  • Handles the non-negativity constraint $\mathbf{w} \ge 0$ natively via box bounds
  • Uses limited-memory BFGS Hessian approximation — O(n·k) memory instead of O(n²)
  • Accepts an analytic gradient (the jac=True flag), making each iteration far faster than finite-difference gradient estimates

Convergence tolerances are set very tight (ftol=1e-12, gtol=1e-8) to ensure a clean solution.


Step 5 — Dose Reconstruction

Once the optimal weights $\mathbf{w}^*$ are found:

$$\mathbf{d}^* = \mathbf{D} \mathbf{w}^*$$

This matrix-vector product reconstructs the dose in all 900 voxels simultaneously. A light Gaussian smooth (σ = 0.8 voxels) is applied purely for visual realism — it mimics the slight dose blurring that occurs in real tissue.


Graph Results

Building dose-influence matrix …
  D_ptv shape: (81, 240)  |  D_oar shape: (28, 240)

Running L-BFGS-B optimisation …
  Converged: False  |  Iterations: 500
  Final objective value: 18031.8827

── Dose Statistics ──────────────────────────────
  PTV   mean=52.53 Gy  min=10.52  max=67.91
  OAR   mean=4.74 Gy  min=0.20  max=10.47
  Normal mean=54.88 Gy

Figure saved → imrt_results.png

Panel-by-Panel Graph Explanation

Panel What it shows
① Phantom Layout The ground truth: where the tumor (PTV) and critical structure (OAR) sit in the patient cross-section.
② Optimised Dose Map The 2-D dose landscape after optimisation. The bright hot-spot aligns tightly with the PTV; the OAR is notably cooler.
③ Beamlet Weight Matrix Each row = one gantry angle; each column = one beamlet. Brighter cells = higher intensity. The optimizer automatically down-weights beamlets that cross the OAR.
④ 3-D Dose Surface The full dose “mountain”. The peak sits over the tumor; the OAR region shows a visible dose shadow — the optimizer has carved out the dose there.
⑤ DVH (Dose-Volume Histogram) The gold standard clinical evaluation plot. Ideal: PTV curve stays high and to the right of the prescription line; OAR curve drops steeply toward low doses.
⑥ Horizontal Dose Profile A 1-D cross-section through the tumour centre. You can directly read the dose fall-off gradient at the PTV–OAR boundary.
⑦ Polar Beam Weight Plot Shows which gantry angles carry the most total intensity. Angles that approach the tumor without crossing the OAR tend to receive higher weights.
⑧ High-Dose Voxel Scatter (3-D) Every voxel receiving ≥ 50% of the prescription dose plotted in 3-D space, coloured by dose. Visually confirms the dose conformity around the PTV.
⑨ Summary Bar Chart Quantitative dose statistics by region for a quick clinical read-out.

Key Clinical Metrics Explained

DVH Target: In a well-optimised plan you want:

The trade-off is controlled entirely by $\lambda$. Increasing $\lambda$ prioritises OAR sparing at the cost of slightly less uniform PTV coverage — a dial the treatment planner can turn to match individual clinical needs.