Perf/adaptfx core hotpath optimizations

adaptive_fractionation_overlap/constants.py

@@ -19,4 +19,12 @@
 
 # Gamma distribution parameters
 DEFAULT_ALPHA = 1.072846744379587
-DEFAULT_BETA = 0.7788684130749829
+DEFAULT_BETA = 0.7788684130749829
+
+# Dynamic programming sentinel values
+INFEASIBLE_VALUE = -1_000_000_000_000.0
+OVERDOSE_STATE_OFFSET = 0.05
+DOSE_GRID_EPSILON = 0.01
+
+COHORT_MAX_OVERLAP_CC = 6.5  # 99th percentile overlap in the 58-patient cohort
+PRECOMPUTE_SCAN_STEP_CC = 0.1

adaptive_fractionation_overlap/helper_functions.py

@@ -3,52 +3,39 @@
 In this file are all helper functions that are needed for the adaptive fractionation calculation
 """
 
-import numpy as np
-from scipy.stats import norm, gamma
+__all__ = [
+    "std_calc",
+    "get_state_space",
+    "probdist",
+    "penalty_calc_single",
+    "penalty_calc_matrix",
+    "max_action",
+    "actual_policy_plotter",
+    "analytic_plotting",
+    "min_dose_to_deliver",
+    "build_dose_decision_lines",
+]
+
+from functools import lru_cache
+
+import matplotlib
 import matplotlib.pyplot as plt
-from .constants import SLOPE, INTERCEPT
-
-
-
-def data_fit(data):
-    """
-    This function fits a normal distribution for the given data
-
-    Parameters
-    ----------
-    data : array or list
-        list with n elements for each observed overlap volume
-
-    Returns
-    -------
-    frozen function
-        normal distribution
-    """
-    mu, std = norm.fit(data)
-    return norm(loc = mu, scale = std)
+import numpy as np
+from scipy.stats import norm
 
-def hyperparam_fit(data):
-    """
-    This function fits the alpha and beta value for the prior
+from .constants import SLOPE, INTERCEPT
 
-    Parameters
-    ----------
-    data : array
-        a nxk matrix with n the amount of patints and k the amount of sparing factors per patient.
+_STD_VALUES = np.arange(0.001, 10, 0.001)
+_STD_VALUES_SQUARED = _STD_VALUES**2
 
-    Returns
-    -------
-    list
-        alpha and beta hyperparameter.
-    """
-    vars = data.var(axis=1)
-    alpha, loc, beta = gamma.fit(vars, floc=0)
-    return [alpha, beta]
+@lru_cache(maxsize=64)
+def _std_prior_term(alpha, beta, n):
+    return _STD_VALUES ** (alpha - n) * np.exp(-_STD_VALUES / beta)
 
-def std_calc(measured_data, alpha, beta):
+def std_calc(measured_data, alpha, beta):
     """
     calculates the most likely standard deviation for a list of k overlap volumes and a gamma prior
-    measured_data: list/array with k sparing factors
+    measured_data: list/array with k overlap volumes
 
     Parameters
     ----------
@@ -57,67 +44,63 @@ def std_calc(measured_data, alpha, beta):
     alpha : float
         shape of gamma distribution
     beta : float
-        scale of gamma distrinbution
+        scale of gamma distribution
 
     Returns
     -------
     std : float
         most likely std based on the measured data and gamma prior
 
     """
-    n = len(measured_data)
-    std_values = np.arange(0.001, 10, 0.001)
-    measured_variance = np.var(measured_data)
-    likelihood_values = (
-        std_values ** (alpha - 1)
-        / std_values ** (n - 1)
-        * np.exp(-1 / beta * std_values)
-        * np.exp(-measured_variance / (2 * (std_values**2 / n)))
-    )
-    std = std_values[np.argmax(likelihood_values)]
-    return std
-
-
+    n = len(measured_data)
+    measured_variance = np.var(measured_data)
+    likelihood_values = _std_prior_term(alpha, beta, n) * np.exp(
+        -measured_variance / (2 * (_STD_VALUES_SQUARED / n))
+    )
+    std = _STD_VALUES[np.argmax(likelihood_values)]
+    return std
 
 def get_state_space(distribution):
     """
     This function spans the state space for different volumes based on a probability distribution
 
     Parameters
     ----------
-    distribution : frozen function
-        normal distribution
+    distribution : tuple(float, float)
+        (mean, std) parameters of a normal distribution
 
     Returns
     -------
-    state_space: Array spanning from the 2% percentile to the 98% percentile with a normalized spacing to define 100 states
+    state_space: Array spanning from the 0.1% percentile to the 99.9% percentile with a normalized spacing to define 200 states
         np.array
     """
-    lower_bound = distribution.ppf(0.001)
-    upper_bound = distribution.ppf(0.999)
+    mean_volume, std_volume = distribution
+    lower_bound = norm.ppf(0.001, loc=mean_volume, scale=std_volume)
+    upper_bound = norm.ppf(0.999, loc=mean_volume, scale=std_volume)
 
     return np.linspace(lower_bound,upper_bound,200)
 
-def probdist(X,state_space):
+def probdist(X,state_space):
     """
     This function produces a probability distribution based on the normal distribution X
 
     Parameters
     ----------
-    X : scipy.stats._distn_infrastructure.rv_frozen
-        distribution function.
+    X : tuple(float, float)
+        (mean, std) parameters of a normal distribution.
 
     Returns
     -------
     prob : np.array
-        array with probabilities for each sparing factor.
+        array with probabilities for each overlap volume.
 
     """
-    spacing = state_space[1]-state_space[0]
-    upper_bounds = state_space + spacing/2
-    lower_bounds = state_space - spacing/2
-    prob = X.cdf(upper_bounds) - X.cdf(lower_bounds)
-    return np.array(prob) #note: this will only add up to roughly 96% instead of 100%
+    spacing = state_space[1]-state_space[0]
+    upper_bounds = state_space + spacing/2
+    lower_bounds = state_space - spacing/2
+    mean_volume, std_volume = X
+    prob = norm.cdf(upper_bounds, loc=mean_volume, scale=std_volume) - norm.cdf(lower_bounds, loc=mean_volume, scale=std_volume)
+    return prob #note: sum depends on how much of the distribution falls within state_space
 
 def penalty_calc_single(physical_dose, min_dose, actual_volume, intercept=INTERCEPT, slope=SLOPE):
     """
@@ -207,8 +190,8 @@ def max_action(accumulated_dose, dose_space, goal):
         gives the size of the resized actionspace to reach the prescribed tumor dose.
 
     """
-    max_action = min(max(dose_space), goal - accumulated_dose)
-    sizer = np.argmin(np.abs(dose_space - max_action))
+    max_deliverable = min(max(dose_space), goal - accumulated_dose)
+    sizer = np.argmin(np.abs(dose_space - max_deliverable))
     sizer = 1 if sizer == 0 else sizer #Make sure that at least the minimum dose is delivered
     return sizer
 
@@ -252,10 +235,12 @@ def analytic_plotting(fraction: int, number_of_fractions: int, values: np.ndarra
     Returns:
         matplotlib.fig: returns a figure with all values plotted as subfigures
     """
-    values[values < -10000000000] = 10000000000
-    min_Value = np.min(values)
-    values[values == 10000000000] = 1.1*min_Value
-    colormap = plt.cm.get_cmap('jet')
+    values = values.copy()
+    plot_infeasible_value = 1e10
+    values[values < -plot_infeasible_value] = plot_infeasible_value
+    min_value = np.min(values)
+    values[values == plot_infeasible_value] = 1.1 * min_value
+    colormap = matplotlib.colormaps['jet']
     number_of_plots = number_of_fractions - fraction
     fig, axs = plt.subplots(1,number_of_plots, figsize = (number_of_plots*10,10))
     if number_of_plots > 1:
@@ -278,7 +263,7 @@ def analytic_plotting(fraction: int, number_of_fractions: int, values: np.ndarra
 
     return fig
 
-def min_dose_to_deliver(accumulated_dose: float, fractions_left: int, prescribed_dose: float, min_dose: float, max_dose: float = None) -> float:
+def min_dose_to_deliver(accumulated_dose: float, fractions_left: int, prescribed_dose: float, min_dose: float, max_dose: float) -> float:
     """
     This function calculates the minimal dose that needs to be delivered in the current fraction to still reach the goal
 
@@ -290,8 +275,8 @@ def min_dose_to_deliver(accumulated_dose: float, fractions_left: int, prescribed
         number of fractions left including the current one
     min_dose : float
         minimal dose that can be delivered in one fraction
-    max_dose : float, optional
-        maximal dose that can be delivered in one fraction, by default None
+    max_dose : float
+        maximal dose that can be delivered in one fraction
 
     Returns
     -------
@@ -339,4 +324,4 @@ def build_dose_decision_lines(volume_x_dose) -> list:
     lines.append(
         f"- Volume ≥ {_format_number(start_vol)} cc: deliver {_format_number(dose)} Gy"
     )
-    return lines
+    return lines

app.py

@@ -10,10 +10,11 @@
     DEFAULT_MIN_DOSE, 
     DEFAULT_MAX_DOSE, 
     DEFAULT_MEAN_DOSE,
-    DEFAULT_DOSE_STEPS, 
+    DEFAULT_DOSE_STEPS,
     DEFAULT_NUMBER_OF_FRACTIONS,
     DEFAULT_ALPHA,
-    DEFAULT_BETA
+    DEFAULT_BETA,
+    INFEASIBLE_VALUE,
 )
 from adaptive_fractionation_overlap.helper_functions import std_calc, build_dose_decision_lines
 
@@ -85,12 +86,12 @@ def build_input_summary(
     return ("\n".join(lines)).encode("utf-8")
 
 @st.cache_data
-def build_precompute_zip(csv_bytes, input_summary_bytes):
+def build_precompute_zip(csv_bytes, input_summary_bytes, fraction):
     """Bundle precompute outputs into a single zip file."""
     buffer = io.BytesIO()
     with zipfile.ZipFile(buffer, mode="w", compression=zipfile.ZIP_DEFLATED) as zf:
         zf.writestr("precomputed_plans.csv", csv_bytes)
-        zf.writestr(f"Fraction_{actual_fraction}_summary.txt", input_summary_bytes)
+        zf.writestr(f"Fraction_{fraction}_summary.txt", input_summary_bytes)
     return buffer.getvalue()
 
 
@@ -116,10 +117,10 @@ def build_precompute_zip(csv_bytes, input_summary_bytes):
         [policies, policies_overlap, volume_space, physical_dose, penalty_added, values, dose_space, probabilities, final_penalty] = af.adaptive_fractionation_core(fraction = int(actual_fraction),volumes = np.array(overlaps), accumulated_dose = float(accumulated_dose), number_of_fractions = int(fractions), min_dose = float(minimum_dose), max_dose = float(maximum_dose), mean_dose = float(mean_dose), dose_steps = float(dose_steps))
         left2, right2 = st.columns(2)
         with left2:
-            actual_value = 'Goal can not be reached' if final_penalty <= -100000000000 else str(np.round(final_penalty,1)) + 'ccGy'
+            actual_value = 'Goal can not be reached' if final_penalty <= INFEASIBLE_VALUE else str(np.round(final_penalty,1)) + 'ccGy'
             st.metric(label="optimal dose for actual fraction", value= str(physical_dose) + 'Gy', delta = (physical_dose - float(mean_dose)))
             st.metric(label="expected final penalty from this fraction", value = actual_value)
-            if final_penalty <= -100000000000:
+            if final_penalty <= INFEASIBLE_VALUE:
                 st.write('the minimal dose is delivered if we overdose, the maximal dose is delivered if we underdose')
                 st.markdown('by taking this approach and delivering the minimum/maximum dose in each fraction we miss the goal by:')
                 st.metric(label= '', value = str(float(accumulated_dose) + float(physical_dose)*(int(fractions) - int(actual_fraction) + 1) - float(mean_dose) * int(fractions)))
@@ -132,7 +133,7 @@ def build_precompute_zip(csv_bytes, input_summary_bytes):
                 st.pyplot(figure)
                 st.write('The figures above show the value function for each future fraction. These functions help to identify whether a potential mistake has been made in the calculation.')
     elif function == 'precompute plan':
-        with st.spinner('computing plans. This might take up to 2-3 minutes'):
+        with st.spinner('computing plans...'):
             volume_x_dose, volumes_to_check, predicted_policies = af.precompute_plan(fraction = int(actual_fraction), volumes = np.array(overlaps), accumulated_dose = float(accumulated_dose), number_of_fractions = int(fractions), min_dose = float(minimum_dose), max_dose = float(maximum_dose), mean_dose = float(mean_dose), dose_steps = float(dose_steps))
         csv = convert_df(volume_x_dose)
         input_summary = build_input_summary(
@@ -151,7 +152,7 @@ def build_precompute_zip(csv_bytes, input_summary_bytes):
             intercept=INTERCEPT,
             volume_x_dose=volume_x_dose
         )
-        zip_bytes = build_precompute_zip(csv, input_summary)
+        zip_bytes = build_precompute_zip(csv, input_summary, int(actual_fraction))
         left2, right2 = st.columns(2)  
         with left2:
             st.dataframe(data = volume_x_dose,height = 600, hide_index = True)
@@ -170,10 +171,10 @@ def build_precompute_zip(csv_bytes, input_summary_bytes):
         cols = st.columns(int(fractions))
         for i, col in enumerate(cols):
             with col:
-                st.metric(label=f"**overlap**", value=f"{overlaps[-(int(fractions) - i)]}cc")
+                st.metric(label="**overlap**", value=f"{overlaps[-(int(fractions) - i)]}cc")
                 st.metric(label=f"**fraction {i + 1}**", value=f"{physical_doses[i]}Gy", delta=np.round(physical_doses[i] - float(mean_dose),1))
         st.header('Plan summary')
         st.markdown('The adaptive plan achieved a total penalty of:')
-        st.metric(label = "penalty", value = str(total_penalty) + 'ccGy', delta = np.round(total_penalty + af.penalty_calc_single(float(mean_dose),6,np.array(overlaps[-int(fractions):]), intercept = INTERCEPT, slope = SLOPE).sum(),2))
+        st.metric(label = "penalty", value = str(total_penalty) + 'ccGy', delta = np.round(total_penalty + af.penalty_calc_single(float(mean_dose),float(minimum_dose),np.array(overlaps[-int(fractions):]), intercept = INTERCEPT, slope = SLOPE).sum(),2))
         st.markdown('The arrow shows the comparison to standard fractionation, i.e. (number of fractions x mean dose). A green arrow shows an improvement.')
 

pyproject.toml

@@ -43,6 +43,5 @@ python_functions = ["test_*"]
 addopts = ["-v", "--tb=short"]
 markers = [
     "unit: Unit tests for individual functions",
-    "integration: Integration tests for multiple components", 
-    "slow: Tests that take longer to run"
-]
+    "integration: Integration tests for multiple components"
+]

tests/conftest.py

@@ -167,9 +167,6 @@ def pytest_configure(config):
     config.addinivalue_line(
         "markers", "integration: marks tests as integration tests (slower, multiple components)"
     )
-    config.addinivalue_line(
-        "markers", "slow: marks tests as slow (performance or stress tests)"
-    )
 
 
 def pytest_collection_modifyitems(config, items):
@@ -179,17 +176,13 @@ def pytest_collection_modifyitems(config, items):
     This function runs after pytest collects all tests and can
     automatically add markers based on test names or locations.
     """
-    # Automatically mark slow tests
     for item in items:
-        if "slow" in item.name or "performance" in item.name:
-            item.add_marker(pytest.mark.slow)
-
         # Mark integration tests
         if "integration" in item.name or "full" in item.name or "end_to_end" in item.name:
             item.add_marker(pytest.mark.integration)
 
         # Mark unit tests (default for most tests)
-        elif not any(marker.name in ["integration", "slow"] for marker in item.iter_markers()):
+        elif not any(marker.name in ["integration"] for marker in item.iter_markers()):
             item.add_marker(pytest.mark.unit)
 
 
@@ -262,14 +255,14 @@ def assert_valid_dose_plan(physical_doses, accumulated_doses, target_dose=None,
     # Check accumulated doses are increasing
     assert_increasing(accumulated_doses)
 
-    # Check accumulated doses are cumulative sums
-    expected_accumulated = np.cumsum(physical_doses)
+    # Check accumulated doses are pre-fraction cumulative sums
+    expected_accumulated = np.concatenate([[0.0], np.cumsum(physical_doses[:-1])])
     assert np.allclose(accumulated_doses, expected_accumulated, atol=1e-10), \
-        "Accumulated doses should be cumulative sum of physical doses"
+        "Accumulated doses should follow algorithm pattern: [0, cumsum(doses[:-1])]"
 
     # Check final dose is reasonable
     if target_dose is not None:
-        final_dose = accumulated_doses[-1]
+        final_dose = accumulated_doses[-1] + physical_doses[-1]
         assert abs(final_dose - target_dose) < tolerance, \
             f"Final dose {final_dose:.1f} should be within {tolerance} Gy of target {target_dose}"
 
@@ -302,4 +295,4 @@ def assert_algorithm_output(result, expected_length=9):
 def generate_random_volumes(n_fractions=5, min_vol=0.0, max_vol=10.0, seed=42):
     """Generate random but reproducible volume data for testing."""
     np.random.seed(seed)
-    return np.random.uniform(min_vol, max_vol, n_fractions + 1)
+    return np.random.uniform(min_vol, max_vol, n_fractions + 1)

tests/test_core_adaptfx.py

@@ -679,7 +679,6 @@ def test_policy_calc_golden_case(self):
         )
         np.testing.assert_allclose(transitions, expected_transitions, atol=1e-12)
 
-
 # Performance and edge case tests
 @pytest.mark.unit
 class TestCoreAdaptfxEdgeCases:
@@ -740,7 +739,6 @@ def test_adaptive_fractionation_core_low_accumulated_dose(self, sample_volumes):
         assert physical_dose == 10.0, "Dose should be maximum dose" 
 
 
-@pytest.mark.slow
 class TestCoreAdaptfxPerformance:
     """Performance tests for core functions."""