- Added py.typed. This makes the package PEP 561 compliant, allowing... (c6a59951) · Commits · Gaetano Scandariato / PyAETNA

PyAETNA/py.typed

0 → 100644

+0 −0

Empty file added.

pyproject.toml

+1 −0

Changes for pyproject.toml: 1 added line, 0 removed lines.
Original line number	Diff line number	Diff line
		@@ -162,6 +162,7 @@ mypy_path = "."
		explicit_package_bases = true
		# Lenient first pass: scientific libraries lack stubs, JAX types are complex.
		# Tighten these gradually as annotations are completed.
		disallow_untyped_defs = true
		ignore_missing_imports = true
		no_implicit_optional = true
		warn_redundant_casts = true

tests/reporting/conftest.py

+3 −2

Changes for tests/reporting/conftest.py: 3 added lines, 2 removed lines.
Original line number	Diff line number	Diff line
		@@ -2,7 +2,6 @@ from dataclasses import dataclass
		from typing import Any

		import numpy as np
		import pytest


		@dataclass
		@@ -26,7 +25,9 @@ class MockPlModel:
		data_types: dict[int, str] \| None = None


		def create_synthetic_reporting_data(n_pts=100, n_visits=5):
		def create_synthetic_reporting_data(
		n_pts: int = 100, n_visits: int = 5
		) -> tuple[MockPlData, MockPlModel, dict[str, Any]]:
		"""
		Creates a minimal synthetic dataset for testing reporting functions.
		"""

tests/test_gap_detection.py

+42 −40

Changes for tests/test_gap_detection.py: 42 added lines, 40 removed lines.
Original line number	Diff line number	Diff line
		@@ -33,7 +33,9 @@ from PyAETNA.reporting.grid_builder import (
		# ─────────────────────────────────────────────────────────────────────────────


		def _two_visit_data(n_per_visit=20, gap_days=10.0):
		def _two_visit_data(
		n_per_visit: int = 20, gap_days: float = 10.0
		) -> tuple[np.ndarray, np.ndarray]:
		"""Return (t_sorted, visit_sorted) for two tight visits separated by a large gap."""
		t1 = np.linspace(0.0, 0.5, n_per_visit)
		t2 = np.linspace(gap_days, gap_days + 0.5, n_per_visit)
		@@ -44,7 +46,7 @@ def _two_visit_data(n_per_visit=20, gap_days=10.0):
		return t, v


		def _three_visit_data(n=10, gap=5.0):
		def _three_visit_data(n: int = 10, gap: float = 5.0) -> tuple[np.ndarray, np.ndarray]:
		"""Return (t, visit) for three visits."""
		t1 = np.linspace(0.0, 0.2, n)
		t2 = np.linspace(gap, gap + 0.2, n)
		@@ -64,19 +66,19 @@ def _three_visit_data(n=10, gap=5.0):
		class TestBuildLcTrainGrid:
		"""Tests for _build_lc_train_grid."""

		def test_output_shapes_consistent(self):
		def test_output_shapes_consistent(self) -> None:
		"""t_grid and visit_grid must have the same length."""
		t, v = _two_visit_data()
		t_grid, visit_grid, _gap_edges = _build_lc_train_grid(t, v)
		assert len(t_grid) == len(visit_grid)

		def test_output_sorted_in_time(self):
		def test_output_sorted_in_time(self) -> None:
		"""Output grid must be monotonically non-decreasing in time."""
		t, v = _two_visit_data()
		t_grid, _, _ = _build_lc_train_grid(t, v)
		assert np.all(np.diff(t_grid) >= 0), "Grid is not time-sorted"

		def test_gap_edges_for_two_visits(self):
		def test_gap_edges_for_two_visits(self) -> None:
		"""With two visits there should be exactly one inter-visit gap edge."""
		t, v = _two_visit_data(gap_days=10.0)
		_, _, gap_edges = _build_lc_train_grid(t, v)
		@@ -84,13 +86,13 @@ class TestBuildLcTrainGrid:
		t_end, t_start = gap_edges[0]
		assert t_end < t_start

		def test_gap_edges_for_three_visits(self):
		def test_gap_edges_for_three_visits(self) -> None:
		"""With three visits there should be exactly two inter-visit gap edges."""
		t, v = _three_visit_data()
		_, _, gap_edges = _build_lc_train_grid(t, v)
		assert len(gap_edges) == 2

		def test_gap_fills_inserted_between_visits(self):
		def test_gap_fills_inserted_between_visits(self) -> None:
		"""Grid must contain points between the two visits (in the inter-visit gap)."""
		t, v = _two_visit_data(n_per_visit=20, gap_days=10.0)
		t_grid, _, _ = _build_lc_train_grid(t, v)
		@@ -98,7 +100,7 @@ class TestBuildLcTrainGrid:
		in_gap = (t_grid > 0.5) & (t_grid < 10.0)
		assert np.any(in_gap), "No gap-fill points inserted between visits"

		def test_training_points_present_in_output(self):
		def test_training_points_present_in_output(self) -> None:
		"""All original training times must appear in (or very close to) the grid."""
		t, v = _two_visit_data(n_per_visit=10, gap_days=5.0)
		t_grid, _, _ = _build_lc_train_grid(t, v)
		@@ -106,14 +108,14 @@ class TestBuildLcTrainGrid:
		dists = np.abs(t_grid - tp)
		assert np.min(dists) < 1e-10, f"Training point {tp} not found in grid"

		def test_single_visit_no_gap_edges(self):
		def test_single_visit_no_gap_edges(self) -> None:
		"""A single-visit dataset should have zero gap edges."""
		t = np.linspace(0.0, 0.5, 30)
		v = np.ones(30, dtype=int)
		_, _, gap_edges = _build_lc_train_grid(t, v)
		assert gap_edges == []

		def test_pad_time_extends_grid(self):
		def test_pad_time_extends_grid(self) -> None:
		"""With pad_time > 0 the grid should extend beyond the data range."""
		t, v = _two_visit_data(n_per_visit=20, gap_days=10.0)
		t_grid_no_pad, _, _ = _build_lc_train_grid(t, v, pad_time=0.0)
		@@ -125,7 +127,7 @@ class TestBuildLcTrainGrid:
		t_grid_padded.max() > t_grid_no_pad.max() + 0.5
		), "Right padding did not extend the grid"

		def test_gap_threshold_controls_in_visit_gap_detection(self):
		def test_gap_threshold_controls_in_visit_gap_detection(self) -> None:
		"""A very large gap_threshold should suppress in-visit gap detection."""
		# Create data with a clear in-visit gap
		t = np.concatenate(
		@@ -143,7 +145,7 @@ class TestBuildLcTrainGrid:
		in_gap
		), "Fine threshold should produce more gap-fill points"

		def test_empty_input_returns_copy(self):
		def test_empty_input_returns_copy(self) -> None:
		"""Empty input should return copies of the empty arrays without crashing."""
		t = np.array([], dtype=float)
		v = np.array([], dtype=int)
		@@ -152,7 +154,7 @@ class TestBuildLcTrainGrid:
		assert len(v_g) == 0
		assert ge == []

		def test_visit_labels_valid_in_output(self):
		def test_visit_labels_valid_in_output(self) -> None:
		"""Visit labels in grid must all come from the original visit set."""
		t, v = _three_visit_data()
		_, visit_grid, _ = _build_lc_train_grid(t, v)
		@@ -168,41 +170,41 @@ class TestBuildLcTrainGrid:
		class TestBuildVisitTimeGrid:
		"""Tests for _build_visit_time_grid (used for RV data)."""

		def test_output_shapes_consistent(self):
		def test_output_shapes_consistent(self) -> None:
		t, v = _two_visit_data()
		t_grid, visit_grid, _ = _build_visit_time_grid(t, v)
		assert len(t_grid) == len(visit_grid)

		def test_output_is_not_empty(self):
		def test_output_is_not_empty(self) -> None:
		t, v = _two_visit_data()
		t_grid, _, _ = _build_visit_time_grid(t, v)
		assert len(t_grid) > 0

		def test_output_sorted_in_time(self):
		def test_output_sorted_in_time(self) -> None:
		t, v = _two_visit_data()
		t_grid, _, _ = _build_visit_time_grid(t, v)
		assert np.all(np.diff(t_grid) >= -1e-12), "Grid not time-sorted"

		def test_gap_edges_two_visits(self):
		def test_gap_edges_two_visits(self) -> None:
		t, v = _two_visit_data(gap_days=10.0)
		_, _, gap_edges = _build_visit_time_grid(t, v)
		assert len(gap_edges) == 1

		def test_gap_fills_in_inter_visit_region(self):
		def test_gap_fills_in_inter_visit_region(self) -> None:
		"""Grid must have points inside the inter-visit gap."""
		t, v = _two_visit_data(n_per_visit=20, gap_days=10.0)
		t_grid, _, _ = _build_visit_time_grid(t, v)
		in_gap = (t_grid > 0.6) & (t_grid < 9.9)
		assert np.any(in_gap), "No gap-fill points inserted between visits"

		def test_n_grid_total_respected_approximately(self):
		def test_n_grid_total_respected_approximately(self) -> None:
		"""Total grid size should be at least n_grid_total (may be slightly larger)."""
		t, v = _two_visit_data(n_per_visit=20, gap_days=8.0)
		t_grid, _, _ = _build_visit_time_grid(t, v, n_grid_total=100)
		# Grid may have more points than n_grid_total due to gap fills
		assert len(t_grid) >= 100, f"Grid size {len(t_grid)} < target 100"

		def test_single_visit(self):
		def test_single_visit(self) -> None:
		"""Single-visit datasets should produce a non-empty grid and no gap edges."""
		t = np.linspace(0.0, 1.0, 20)
		v = np.ones(20, dtype=int)
		@@ -210,14 +212,14 @@ class TestBuildVisitTimeGrid:
		assert len(t_grid) > 0
		assert gap_edges == []

		def test_pad_time_extends_grid_beyond_data(self):
		def test_pad_time_extends_grid_beyond_data(self) -> None:
		t, v = _two_visit_data(n_per_visit=10, gap_days=5.0)
		t_grid_nopad, _, _ = _build_visit_time_grid(t, v, pad_time=0.0)
		t_grid_padded, _, _ = _build_visit_time_grid(t, v, pad_time=0.5)
		assert t_grid_padded.min() < t_grid_nopad.min()
		assert t_grid_padded.max() > t_grid_nopad.max()

		def test_min_pts_per_visit_enforced(self):
		def test_min_pts_per_visit_enforced(self) -> None:
		"""Each visit must have at least min_pts_per_visit grid points."""
		t, v = _two_visit_data(n_per_visit=5, gap_days=8.0)
		_t_grid, visit_grid, _ = _build_visit_time_grid(
		@@ -227,7 +229,7 @@ class TestBuildVisitTimeGrid:
		pts = np.sum(visit_grid == uv)
		assert pts >= 15, f"Visit {uv} has only {pts} grid pts (< 15)"

		def test_visit_labels_valid(self):
		def test_visit_labels_valid(self) -> None:
		t, v = _three_visit_data()
		_, visit_grid, _ = _build_visit_time_grid(t, v)
		allowed = set(np.unique(v).tolist())
		@@ -243,7 +245,7 @@ class TestBuildVisitTimeGrid:
		class TestBuildAugmentedX:
		"""Tests for _build_augmented_x (gap-fill pseudo-observation builder)."""

		def test_output_length_is_train_plus_pseudo(self):
		def test_output_length_is_train_plus_pseudo(self) -> None:
		"""x_aug length == N_train + N_pseudo where N_pseudo <= len(x_grid)."""
		x_train = np.linspace(0.0, 1.0, 10)
		x_grid = np.linspace(0.5, 2.0, 20) # some overlap, some new
		@@ -251,13 +253,13 @@ class TestBuildAugmentedX:
		n_pseudo = int(np.sum(~is_real))
		assert len(x_aug) == len(x_train) + n_pseudo

		def test_x_aug_is_sorted(self):
		def test_x_aug_is_sorted(self) -> None:
		x_train = np.linspace(0.0, 1.0, 15)
		x_grid = np.linspace(1.5, 3.0, 20)
		x_aug, _, _, _ = _build_augmented_x(x_train, x_grid)
		assert np.all(np.diff(x_aug) >= 0), "x_aug is not sorted"

		def test_is_real_marks_training_points(self):
		def test_is_real_marks_training_points(self) -> None:
		"""Positions matching training points must be marked as real."""
		x_train = np.array([0.0, 1.0, 2.0, 3.0])
		x_grid = np.array([0.0, 1.0, 2.0, 3.0, 4.0]) # 4.0 is new
		@@ -267,7 +269,7 @@ class TestBuildAugmentedX:
		idx = np.argmin(np.abs(x_aug - xt))
		assert is_real[idx], f"Training point {xt} not marked as real"

		def test_real_order_maps_training_to_aug_positions(self):
		def test_real_order_maps_training_to_aug_positions(self) -> None:
		"""real_order[i] must be the index of x_train[i] in x_aug."""
		x_train = np.array([0.0, 2.0, 4.0])
		x_grid = np.array([1.0, 3.0, 5.0])
		@@ -277,7 +279,7 @@ class TestBuildAugmentedX:
		abs(x_aug[real_order[i]] - xt) < 1e-9
		), f"real_order[{i}]={real_order[i]} does not map to x_train[{i}]={xt}"

		def test_aug_to_grid_maps_pseudo_to_grid_indices(self):
		def test_aug_to_grid_maps_pseudo_to_grid_indices(self) -> None:
		"""aug_to_grid values must be valid indices into x_grid."""
		x_train = np.linspace(0.0, 0.5, 10)
		x_grid = np.linspace(1.0, 2.0, 8) # all pseudo (no overlap)
		@@ -287,7 +289,7 @@ class TestBuildAugmentedX:
		assert np.all(aug_to_grid >= 0)
		assert np.all(aug_to_grid < len(x_grid))

		def test_overlapping_grid_points_not_duplicated(self):
		def test_overlapping_grid_points_not_duplicated(self) -> None:
		"""Grid points coinciding with training points should not add pseudo-obs."""
		x_train = np.linspace(0.0, 1.0, 11) # 0.0, 0.1, ..., 1.0
		x_grid = np.array([0.0, 0.1, 0.2]) # all coincide with training
		@@ -297,7 +299,7 @@ class TestBuildAugmentedX:
		len(aug_to_grid) == 0
		), f"Expected 0 pseudo-obs for fully overlapping grid, got {len(aug_to_grid)}"

		def test_empty_grid_returns_only_training(self):
		def test_empty_grid_returns_only_training(self) -> None:
		x_train = np.linspace(0.0, 1.0, 5)
		x_grid = np.array([])
		x_aug, is_real, _real_order, aug_to_grid = _build_augmented_x(x_train, x_grid)
		@@ -305,7 +307,7 @@ class TestBuildAugmentedX:
		assert np.all(is_real)
		assert len(aug_to_grid) == 0

		def test_empty_training_returns_all_pseudo(self):
		def test_empty_training_returns_all_pseudo(self) -> None:
		x_train = np.array([])
		x_grid = np.linspace(0.0, 1.0, 5)
		x_aug, is_real, _real_order, aug_to_grid = _build_augmented_x(x_train, x_grid)
		@@ -313,7 +315,7 @@ class TestBuildAugmentedX:
		assert np.all(~is_real)
		assert len(aug_to_grid) == len(x_grid)

		def test_aug_to_grid_covers_all_pseudo_unique(self):
		def test_aug_to_grid_covers_all_pseudo_unique(self) -> None:
		"""Each pseudo-observation should map to a distinct x_grid index."""
		x_train = np.linspace(0.0, 0.5, 5)
		x_grid = np.linspace(2.0, 3.0, 8) # all pseudo
		@@ -330,32 +332,32 @@ class TestBuildAugmentedX:
		class TestGridBuilderClass:
		"""Tests for the GridBuilder class in grid_builder.py."""

		def test_lc_strategy_returns_valid_grid(self):
		def test_lc_strategy_returns_valid_grid(self) -> None:
		t, v = _two_visit_data()
		gb = GridBuilder(t, v)
		t_grid, visit_grid, _gap_edges = gb.build_time_grid(strategy="lc")
		assert len(t_grid) > 0
		assert len(t_grid) == len(visit_grid)

		def test_visit_strategy_returns_valid_grid(self):
		def test_visit_strategy_returns_valid_grid(self) -> None:
		t, v = _two_visit_data()
		gb = GridBuilder(t, v)
		t_grid, visit_grid, _gap_edges = gb.build_time_grid(strategy="visit")
		assert len(t_grid) > 0
		assert len(t_grid) == len(visit_grid)

		def test_invalid_strategy_raises(self):
		def test_invalid_strategy_raises(self) -> None:
		t, v = _two_visit_data()
		gb = GridBuilder(t, v)
		with pytest.raises(ValueError, match="Unknown grid strategy"):
		gb.build_time_grid(strategy="invalid")

		def test_unique_visits_stored(self):
		def test_unique_visits_stored(self) -> None:
		t, v = _three_visit_data()
		gb = GridBuilder(t, v)
		assert len(gb.unique_visits) == 3

		def test_lc_strategy_has_gap_fills(self):
		def test_lc_strategy_has_gap_fills(self) -> None:
		t, v = _two_visit_data(gap_days=8.0)
		gb = GridBuilder(t, v)
		t_grid, _, _ = gb.build_time_grid(strategy="lc")
		@@ -371,18 +373,18 @@ class TestGridBuilderClass:
		class TestBuildTimeGridFunctional:
		"""Tests for the build_time_grid() module-level function."""

		def test_lc_strategy(self):
		def test_lc_strategy(self) -> None:
		t, v = _two_visit_data()
		t_grid, visit_grid, _ = build_time_grid(t, v, strategy="lc")
		assert len(t_grid) > 0
		assert len(t_grid) == len(visit_grid)

		def test_visit_strategy(self):
		def test_visit_strategy(self) -> None:
		t, v = _two_visit_data()
		t_grid, _visit_grid, _ = build_time_grid(t, v, strategy="visit")
		assert len(t_grid) > 0

		def test_kwargs_forwarded(self):
		def test_kwargs_forwarded(self) -> None:
		t, v = _two_visit_data()
		t_small, _, _ = build_time_grid(
		t, v, strategy="visit", n_grid_total=50, min_pts_per_visit=10