oscura 0.5.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

oscura/inference/__init__.py

@@ -123,6 +123,8 @@ from oscura.inference.spectral import auto_spectral_config
 from oscura.inference.state_machine import (
     FiniteAutomaton,
     State,
+    StateMachine,
+    StateMachineExtractor,
     StateMachineInferrer,
     Transition,
     infer_rpni,
@@ -196,6 +198,8 @@ __all__ = [
     "SequentialBayesian",
     "SimulatorTeacher",
     "State",
+    "StateMachine",
+    "StateMachineExtractor",
     "StateMachineInferrer",
     "StreamSegment",
     "TCPStreamReassembler",

oscura/inference/active_learning/observation_table.py

@@ -219,10 +219,13 @@ class ObservationTable:
         initial_row = self.row(())
         initial_state = row_to_state[initial_row]
 
+        # Type cast alphabet from set[str] to set[str | int] for FiniteAutomaton
+        alphabet_union: set[str | int] = set(self.alphabet)
+
         return FiniteAutomaton(
             states=states,
             transitions=transitions,
-            alphabet=self.alphabet.copy(),
+            alphabet=alphabet_union,
             initial_state=initial_state,
             accepting_states=accepting_states,
         )

oscura/inference/alignment.py

@@ -122,42 +122,12 @@ def align_global(
                 traceback[i, j] = 2  # Left
 
     # Traceback to get alignment
-    aligned_a = []
-    aligned_b = []
-
-    i, j = n, m
-    while i > 0 or j > 0:
-        if traceback[i, j] == 0:  # Diagonal
-            aligned_a.append(int(arr_a[i - 1]))
-            aligned_b.append(int(arr_b[j - 1]))
-            i -= 1
-            j -= 1
-        elif traceback[i, j] == 1:  # Up
-            aligned_a.append(int(arr_a[i - 1]))
-            aligned_b.append(-1)  # Gap
-            i -= 1
-        else:  # Left
-            aligned_a.append(-1)  # Gap
-            aligned_b.append(int(arr_b[j - 1]))
-            j -= 1
-
-    # Reverse (we traced backwards)
-    aligned_a = list(reversed(aligned_a))
-    aligned_b = list(reversed(aligned_b))
+    aligned_a, aligned_b = _traceback_alignment(traceback, arr_a, arr_b, n, m)
 
     # Calculate statistics
     final_score = float(score_matrix[n, m])
     similarity = compute_similarity(aligned_a, aligned_b)
-
-    # Handle empty alignments
-    if len(aligned_a) == 0:
-        identity = 0.0
-        gaps = 0
-    else:
-        identity = sum(
-            1 for a, b in zip(aligned_a, aligned_b, strict=True) if a == b and a != -1
-        ) / len(aligned_a)
-        gaps = sum(1 for a, b in zip(aligned_a, aligned_b, strict=True) if a == -1 or b == -1)
+    identity, gaps = _calculate_alignment_stats(aligned_a, aligned_b)
 
     # Find conserved and variable regions
     conserved = _find_conserved_simple(aligned_a, aligned_b)
@@ -196,97 +166,12 @@ def align_local(
     Returns:
         AlignmentResult with best local alignment
     """
-    # Convert to arrays
-    if isinstance(seq_a, bytes):
-        arr_a = np.frombuffer(seq_a, dtype=np.uint8)
-    else:
-        arr_a = np.array(seq_a, dtype=np.uint8)
-
-    if isinstance(seq_b, bytes):
-        arr_b = np.frombuffer(seq_b, dtype=np.uint8)
-    else:
-        arr_b = np.array(seq_b, dtype=np.uint8)
-
-    n, m = len(arr_a), len(arr_b)
-
-    # Initialize scoring matrix and traceback matrix
-    score_matrix = np.zeros((n + 1, m + 1), dtype=np.float32)
-    traceback = np.zeros((n + 1, m + 1), dtype=np.int8)
-
-    # Track maximum score position
-    max_score = 0.0
-    max_i, max_j = 0, 0
-
-    # Fill the matrices (Smith-Waterman: no negative scores)
-    for i in range(1, n + 1):
-        for j in range(1, m + 1):
-            # Match/mismatch
-            if arr_a[i - 1] == arr_b[j - 1]:
-                diag_score = score_matrix[i - 1, j - 1] + match_score
-            else:
-                diag_score = score_matrix[i - 1, j - 1] + mismatch_penalty
-
-            # Gap in seq_b (up)
-            up_score = score_matrix[i - 1, j] + gap_penalty
-
-            # Gap in seq_a (left)
-            left_score = score_matrix[i, j - 1] + gap_penalty
-
-            # Smith-Waterman: can start fresh (score = 0)
-            cell_score = max(0.0, diag_score, up_score, left_score)
-            score_matrix[i, j] = cell_score
-
-            if cell_score == 0:
-                traceback[i, j] = -1  # Stop
-            elif cell_score == diag_score:
-                traceback[i, j] = 0  # Diagonal
-            elif cell_score == up_score:
-                traceback[i, j] = 1  # Up
-            else:
-                traceback[i, j] = 2  # Left
-
-            # Track maximum
-            if cell_score > max_score:
-                max_score = cell_score
-                max_i, max_j = i, j
-
-    # Traceback from max position
-    aligned_a = []
-    aligned_b = []
-
-    i, j = max_i, max_j
-    while i > 0 and j > 0 and traceback[i, j] != -1:
-        if traceback[i, j] == 0:  # Diagonal
-            aligned_a.append(int(arr_a[i - 1]))
-            aligned_b.append(int(arr_b[j - 1]))
-            i -= 1
-            j -= 1
-        elif traceback[i, j] == 1:  # Up
-            aligned_a.append(int(arr_a[i - 1]))
-            aligned_b.append(-1)  # Gap
-            i -= 1
-        else:  # Left
-            aligned_a.append(-1)  # Gap
-            aligned_b.append(int(arr_b[j - 1]))
-            j -= 1
-
-    # Reverse
-    aligned_a = list(reversed(aligned_a))
-    aligned_b = list(reversed(aligned_b))
-
-    # Calculate statistics
-    if len(aligned_a) > 0:
-        similarity = compute_similarity(aligned_a, aligned_b)
-        identity = sum(
-            1 for a, b in zip(aligned_a, aligned_b, strict=True) if a == b and a != -1
-        ) / len(aligned_a)
-        gaps = sum(1 for a, b in zip(aligned_a, aligned_b, strict=True) if a == -1 or b == -1)
-    else:
-        similarity = 0.0
-        identity = 0.0
-        gaps = 0
-
-    # Find conserved and variable regions
+    arr_a, arr_b = _convert_to_arrays(seq_a, seq_b)
+    score_matrix, traceback, max_score, max_pos = _build_sw_matrix(
+        arr_a, arr_b, gap_penalty, match_score, mismatch_penalty
+    )
+    aligned_a, aligned_b = _traceback_local(traceback, arr_a, arr_b, max_pos)
+    similarity, identity, gaps = _compute_local_stats(aligned_a, aligned_b)
     conserved = _find_conserved_simple(aligned_a, aligned_b)
     variable = _find_variable_simple(aligned_a, aligned_b)
 
@@ -532,6 +417,214 @@ def find_variable_regions(
     return regions
 
 
+def _convert_to_arrays(
+    seq_a: bytes | NDArray[Any], seq_b: bytes | NDArray[Any]
+) -> tuple[NDArray[np.uint8], NDArray[np.uint8]]:
+    """Convert input sequences to numpy arrays.
+
+    Args:
+        seq_a: First sequence (bytes or array).
+        seq_b: Second sequence (bytes or array).
+
+    Returns:
+        Tuple of (arr_a, arr_b) as uint8 arrays.
+    """
+    if isinstance(seq_a, bytes):
+        arr_a = np.frombuffer(seq_a, dtype=np.uint8)
+    else:
+        arr_a = np.array(seq_a, dtype=np.uint8)
+
+    if isinstance(seq_b, bytes):
+        arr_b = np.frombuffer(seq_b, dtype=np.uint8)
+    else:
+        arr_b = np.array(seq_b, dtype=np.uint8)
+
+    return arr_a, arr_b
+
+
+def _build_sw_matrix(
+    arr_a: NDArray[np.uint8],
+    arr_b: NDArray[np.uint8],
+    gap_penalty: float,
+    match_score: float,
+    mismatch_penalty: float,
+) -> tuple[NDArray[np.float32], NDArray[np.int8], float, tuple[int, int]]:
+    """Build Smith-Waterman scoring and traceback matrices.
+
+    Args:
+        arr_a: First sequence array.
+        arr_b: Second sequence array.
+        gap_penalty: Penalty for gaps.
+        match_score: Score for matches.
+        mismatch_penalty: Penalty for mismatches.
+
+    Returns:
+        (score_matrix, traceback, max_score, max_position).
+    """
+    n, m = len(arr_a), len(arr_b)
+    score_matrix = np.zeros((n + 1, m + 1), dtype=np.float32)
+    traceback = np.zeros((n + 1, m + 1), dtype=np.int8)
+    max_score = 0.0
+    max_i, max_j = 0, 0
+
+    for i in range(1, n + 1):
+        for j in range(1, m + 1):
+            # Match/mismatch
+            if arr_a[i - 1] == arr_b[j - 1]:
+                diag_score = score_matrix[i - 1, j - 1] + match_score
+            else:
+                diag_score = score_matrix[i - 1, j - 1] + mismatch_penalty
+
+            up_score = score_matrix[i - 1, j] + gap_penalty
+            left_score = score_matrix[i, j - 1] + gap_penalty
+
+            # Smith-Waterman: can start fresh (score = 0)
+            cell_score = max(0.0, diag_score, up_score, left_score)
+            score_matrix[i, j] = cell_score
+
+            if cell_score == 0:
+                traceback[i, j] = -1  # Stop
+            elif cell_score == diag_score:
+                traceback[i, j] = 0  # Diagonal
+            elif cell_score == up_score:
+                traceback[i, j] = 1  # Up
+            else:
+                traceback[i, j] = 2  # Left
+
+            if cell_score > max_score:
+                max_score = cell_score
+                max_i, max_j = i, j
+
+    return score_matrix, traceback, max_score, (max_i, max_j)
+
+
+def _traceback_local(
+    traceback: NDArray[np.int8],
+    arr_a: NDArray[np.uint8],
+    arr_b: NDArray[np.uint8],
+    max_pos: tuple[int, int],
+) -> tuple[list[int], list[int]]:
+    """Perform traceback from max position for local alignment.
+
+    Args:
+        traceback: Traceback matrix.
+        arr_a: First sequence array.
+        arr_b: Second sequence array.
+        max_pos: (i, j) position of maximum score.
+
+    Returns:
+        Tuple of (aligned_a, aligned_b) with -1 for gaps.
+    """
+    aligned_a = []
+    aligned_b = []
+
+    i, j = max_pos
+    while i > 0 and j > 0 and traceback[i, j] != -1:
+        if traceback[i, j] == 0:  # Diagonal
+            aligned_a.append(int(arr_a[i - 1]))
+            aligned_b.append(int(arr_b[j - 1]))
+            i -= 1
+            j -= 1
+        elif traceback[i, j] == 1:  # Up
+            aligned_a.append(int(arr_a[i - 1]))
+            aligned_b.append(-1)  # Gap
+            i -= 1
+        else:  # Left
+            aligned_a.append(-1)  # Gap
+            aligned_b.append(int(arr_b[j - 1]))
+            j -= 1
+
+    return list(reversed(aligned_a)), list(reversed(aligned_b))
+
+
+def _compute_local_stats(aligned_a: list[int], aligned_b: list[int]) -> tuple[float, float, int]:
+    """Compute statistics for local alignment.
+
+    Args:
+        aligned_a: First aligned sequence.
+        aligned_b: Second aligned sequence.
+
+    Returns:
+        (similarity, identity, gaps).
+    """
+    if len(aligned_a) == 0:
+        return 0.0, 0.0, 0
+
+    similarity = compute_similarity(aligned_a, aligned_b)
+    identity = sum(
+        1 for a, b in zip(aligned_a, aligned_b, strict=True) if a == b and a != -1
+    ) / len(aligned_a)
+    gaps = sum(1 for a, b in zip(aligned_a, aligned_b, strict=True) if a == -1 or b == -1)
+
+    return similarity, identity, gaps
+
+
+def _traceback_alignment(
+    traceback: NDArray[np.int8],
+    arr_a: NDArray[np.uint8],
+    arr_b: NDArray[np.uint8],
+    n: int,
+    m: int,
+) -> tuple[list[int], list[int]]:
+    """Perform traceback to extract aligned sequences.
+
+    Args:
+        traceback: Traceback matrix (0=diagonal, 1=up, 2=left).
+        arr_a: First sequence array.
+        arr_b: Second sequence array.
+        n: Length of first sequence.
+        m: Length of second sequence.
+
+    Returns:
+        Tuple of (aligned_a, aligned_b) with -1 for gaps.
+    """
+    aligned_a = []
+    aligned_b = []
+
+    i, j = n, m
+    while i > 0 or j > 0:
+        if traceback[i, j] == 0:  # Diagonal (match/mismatch)
+            aligned_a.append(int(arr_a[i - 1]))
+            aligned_b.append(int(arr_b[j - 1]))
+            i -= 1
+            j -= 1
+        elif traceback[i, j] == 1:  # Up (gap in seq_b)
+            aligned_a.append(int(arr_a[i - 1]))
+            aligned_b.append(-1)  # Gap
+            i -= 1
+        else:  # Left (gap in seq_a)
+            aligned_a.append(-1)  # Gap
+            aligned_b.append(int(arr_b[j - 1]))
+            j -= 1
+
+    # Reverse (we traced backwards)
+    return list(reversed(aligned_a)), list(reversed(aligned_b))
+
+
+def _calculate_alignment_stats(
+    aligned_a: list[int],
+    aligned_b: list[int],
+) -> tuple[float, int]:
+    """Calculate identity and gap statistics for alignment.
+
+    Args:
+        aligned_a: First aligned sequence (-1 for gaps).
+        aligned_b: Second aligned sequence (-1 for gaps).
+
+    Returns:
+        Tuple of (identity, gaps).
+    """
+    if len(aligned_a) == 0:
+        return 0.0, 0
+
+    identity = sum(
+        1 for a, b in zip(aligned_a, aligned_b, strict=True) if a == b and a != -1
+    ) / len(aligned_a)
+    gaps = sum(1 for a, b in zip(aligned_a, aligned_b, strict=True) if a == -1 or b == -1)
+
+    return identity, gaps
+
+
 def _find_conserved_simple(aligned_a: list[int], aligned_b: list[int]) -> list[tuple[int, int]]:
     """Find conserved regions in pairwise alignment.
 

oscura/inference/bayesian.py

@@ -142,9 +142,9 @@ class Prior:
             ValueError: If distribution is not recognized.
         """
         if self.distribution == "normal":
-            return float(stats.norm.pdf(x, loc=self.params["mean"], scale=self.params["std"]))  # type: ignore[no-any-return]
+            return float(stats.norm.pdf(x, loc=self.params["mean"], scale=self.params["std"]))
         elif self.distribution == "uniform":
-            return float(  # type: ignore[no-any-return]
+            return float(
                 stats.uniform.pdf(
                     x, loc=self.params["low"], scale=self.params["high"] - self.params["low"]
                 )
@@ -156,16 +156,16 @@ class Prior:
             log_x = np.log(np.maximum(x, 1e-100))  # Avoid log(0)
             density = stats.uniform.pdf(log_x, loc=log_low, scale=log_high - log_low)
             # Jacobian correction: d(log x)/dx = 1/x
-            result = density / np.maximum(x, 1e-100)
-            return result  # type: ignore[return-value, no-any-return]
+            result: float | NDArray[np.floating[Any]] = density / np.maximum(x, 1e-100)
+            return result
         elif self.distribution == "beta":
-            return float(stats.beta.pdf(x, a=self.params["a"], b=self.params["b"]))  # type: ignore[no-any-return]
+            return float(stats.beta.pdf(x, a=self.params["a"], b=self.params["b"]))
         elif self.distribution == "gamma":
-            return float(stats.gamma.pdf(x, a=self.params["shape"], scale=self.params["scale"]))  # type: ignore[no-any-return]
+            return float(stats.gamma.pdf(x, a=self.params["shape"], scale=self.params["scale"]))
         elif self.distribution == "half_normal":
-            return float(stats.halfnorm.pdf(x, scale=self.params["scale"]))  # type: ignore[no-any-return]
+            return float(stats.halfnorm.pdf(x, scale=self.params["scale"]))
         elif self.distribution == "geometric":
-            return float(stats.geom.pmf(x, p=self.params["p"]))  # type: ignore[no-any-return]
+            return float(stats.geom.pmf(x, p=self.params["p"]))
         else:
             raise ValueError(f"PDF not implemented for {self.distribution}")
 
@@ -191,7 +191,7 @@ class Prior:
             # Sample uniformly on log scale, then exponentiate
             log_low = np.log(self.params["low"])
             log_high = np.log(self.params["high"])
-            log_samples = stats.uniform.rvs(loc=log_low, scale=log_high - log_low, size=n)  # type: ignore[no-any-return]
+            log_samples = stats.uniform.rvs(loc=log_low, scale=log_high - log_low, size=n)
             return np.exp(log_samples)  # type: ignore[no-any-return]
         elif self.distribution == "beta":
             return stats.beta.rvs(a=self.params["a"], b=self.params["b"], size=n)  # type: ignore[no-any-return]
@@ -364,7 +364,25 @@ class BayesianInference:
             >>> inference = BayesianInference()
             >>> posterior = inference.update("frequency", likelihood)
         """
-        # Get prior
+        prior = self._get_prior(param, prior)
+        samples = self._sample_from_prior(prior, param, num_samples)
+        likelihoods = self._compute_likelihoods(samples, likelihood_fn, param)
+        weights = self._normalize_weights(likelihoods, param)
+        return self._build_posterior(samples, weights)
+
+    def _get_prior(self, param: str, prior: Prior | None) -> Prior:
+        """Get prior distribution for parameter.
+
+        Args:
+            param: Parameter name.
+            prior: Optional explicit prior.
+
+        Returns:
+            Prior distribution to use.
+
+        Raises:
+            ValueError: If parameter unknown and no prior provided.
+        """
         if prior is None:
             if param not in self.priors:
                 raise ValueError(
@@ -372,17 +390,51 @@ class BayesianInference:
                     f"Known parameters: {list(self.priors.keys())}"
                 )
             prior = self.priors[param]
+        return prior
+
+    def _sample_from_prior(
+        self, prior: Prior, param: str, num_samples: int
+    ) -> NDArray[np.floating[Any]]:
+        """Sample from prior distribution.
 
-        # Sample from prior
+        Args:
+            prior: Prior distribution.
+            param: Parameter name for error messages.
+            num_samples: Number of samples.
+
+        Returns:
+            Array of prior samples.
+
+        Raises:
+            AnalysisError: If sampling fails.
+        """
         try:
-            samples = prior.sample(num_samples)
+            return prior.sample(num_samples)
         except Exception as e:
             raise AnalysisError(
                 f"Failed to sample from prior for '{param}'",
                 details=str(e),
             ) from e
 
-        # Compute likelihood for each sample
+    def _compute_likelihoods(
+        self,
+        samples: NDArray[np.floating[Any]],
+        likelihood_fn: Callable[[float], float],
+        param: str,
+    ) -> NDArray[np.floating[Any]]:
+        """Compute likelihoods for all samples.
+
+        Args:
+            samples: Prior samples.
+            likelihood_fn: Likelihood function.
+            param: Parameter name for error messages.
+
+        Returns:
+            Array of likelihood values.
+
+        Raises:
+            AnalysisError: If likelihood computation fails or all zeros.
+        """
         try:
             likelihoods = np.array([likelihood_fn(s) for s in samples])
         except Exception as e:
@@ -392,7 +444,6 @@ class BayesianInference:
                 fix_hint="Check that likelihood_fn is compatible with prior samples",
             ) from e
 
-        # Check for valid likelihoods (before numerical stability fixes)
         if np.all(likelihoods == 0):
             raise AnalysisError(
                 f"All likelihood values are zero for '{param}'",
@@ -400,42 +451,71 @@ class BayesianInference:
                 fix_hint="Adjust prior range or check likelihood function",
             )
 
-        # Numerical stability: normalize likelihoods to prevent underflow
-        # Use log-space computation if likelihoods are very small
+        return likelihoods
+
+    def _normalize_weights(
+        self, likelihoods: NDArray[np.floating[Any]], param: str
+    ) -> NDArray[np.floating[Any]]:
+        """Normalize likelihoods to importance weights.
+
+        Uses numerical stability techniques for extreme values.
+
+        Args:
+            likelihoods: Likelihood values.
+            param: Parameter name for error messages.
+
+        Returns:
+            Normalized importance weights.
+
+        Raises:
+            AnalysisError: If all likelihoods are zero.
+        """
         max_likelihood = np.max(likelihoods)
-        if max_likelihood > 0:
-            # Normalize by max to prevent overflow/underflow
-            normalized_likelihoods = likelihoods / max_likelihood
-            # Check if we need log-space (very small values)
-            if max_likelihood < 1e-300:
-                # Use log-space for extreme underflow
-                log_likelihoods = np.log(np.maximum(likelihoods, 1e-300))
-                log_likelihoods -= np.max(log_likelihoods)  # Normalize
-                weights = np.exp(log_likelihoods)
-                weights /= np.sum(weights)
-            else:
-                # Standard normalization
-                weights = normalized_likelihoods / np.sum(normalized_likelihoods)
-        else:
-            # All likelihoods are zero - this should have been caught above
+        if max_likelihood <= 0:
             raise AnalysisError(
                 f"All likelihood values are zero for '{param}'",
                 details="Observation may be incompatible with prior range",
                 fix_hint="Adjust prior range or check likelihood function",
             )
 
+        # Normalize by max to prevent overflow/underflow
+        normalized_likelihoods = likelihoods / max_likelihood
+
+        # Use log-space for extreme underflow
+        if max_likelihood < 1e-300:
+            log_likelihoods = np.log(np.maximum(likelihoods, 1e-300))
+            log_likelihoods -= np.max(log_likelihoods)
+            weights = np.exp(log_likelihoods)
+            weights /= np.sum(weights)
+        else:
+            weights = normalized_likelihoods / np.sum(normalized_likelihoods)
+
+        result: NDArray[np.floating[Any]] = np.asarray(weights, dtype=np.float64)
+        return result
+
+    def _build_posterior(
+        self, samples: NDArray[np.floating[Any]], weights: NDArray[np.floating[Any]]
+    ) -> Posterior:
+        """Build posterior distribution from weighted samples.
+
+        Args:
+            samples: Prior samples.
+            weights: Importance weights.
+
+        Returns:
+            Posterior distribution with statistics.
+        """
         # Compute posterior statistics
         mean = float(np.sum(samples * weights))
         variance = float(np.sum(weights * (samples - mean) ** 2))
         std = float(np.sqrt(variance))
 
-        # Compute 95% credible interval via weighted percentiles
+        # Compute 95% credible interval
         sorted_indices = np.argsort(samples)
         sorted_samples = samples[sorted_indices]
         sorted_weights = weights[sorted_indices]
         cumsum = np.cumsum(sorted_weights)
 
-        # Find 2.5th and 97.5th percentiles
         ci_lower = float(sorted_samples[np.searchsorted(cumsum, 0.025)])
         ci_upper = float(sorted_samples[np.searchsorted(cumsum, 0.975)])