goad-py 0.6.0__cp38-abi3-musllinux_1_2_aarch64.whl → 0.7.0__cp38-abi3-musllinux_1_2_aarch64.whl

goad_py/convergence.py

@@ -1,10 +1,20 @@
-from dataclasses import dataclass
-from typing import List, Dict, Optional, Tuple
-import numpy as np
-from . import _goad_py as goad
+import contextlib
 import os
 import random
+import sys
+from dataclasses import dataclass
 from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+from rich.console import Console
+
+from . import _goad_py as goad
+from .convergence_display import (
+    ArrayConvergenceVariable,
+    ConvergenceDisplay,
+    ConvergenceVariable,
+)
 
 
 @dataclass
@@ -96,6 +106,8 @@ class Convergence:
             mueller_2d: Whether to collect 2D Mueller matrices
         """
         self.settings = settings
+        # Enable quiet mode to suppress Rust progress bars
+        self.settings.quiet = True
         self.convergables = convergables
         self.batch_size = batch_size
         self.max_orientations = max_orientations
@@ -129,6 +141,39 @@ class Convergence:
         self.mueller_1d_sum = None
         self.mueller_2d_sum = None
 
+        # Rich console
+        self._console = Console()
+
+        # Create display variables for the new display system
+        display_variables = []
+        for conv in self.convergables:
+            if conv.is_mueller():
+                display_variables.append(
+                    ArrayConvergenceVariable(
+                        name=conv.variable,
+                        tolerance=conv.tolerance,
+                        tolerance_type=conv.tolerance_type,
+                        indices=conv.theta_indices,
+                    )
+                )
+            else:
+                display_variables.append(
+                    ConvergenceVariable(
+                        name=conv.variable,
+                        tolerance=conv.tolerance,
+                        tolerance_type=conv.tolerance_type,
+                    )
+                )
+
+        # Initialize display system
+        self._display = ConvergenceDisplay(
+            variables=display_variables,
+            batch_size=self.batch_size,
+            min_batches=self.min_batches,
+            convergence_type=self._get_convergence_type(),
+            console=self._console,
+        )
+
     def _update_statistics(self, results: goad.Results, batch_size: int):
         """Update statistics with new batch results.
 
@@ -425,144 +470,84 @@ class Convergence:
         converged_status = self._check_convergence()
         return all(converged_status.values())
 
-    def _print_progress(self, iteration: int):
-        """Print convergence progress.
+    def _get_convergence_type(self) -> str:
+        """Get the convergence type name for display."""
+        class_name = self.__class__.__name__
+        if class_name == "EnsembleConvergence":
+            return "Ensemble"
+        elif class_name == "Convergence":
+            return "Standard"
+        else:
+            return class_name
+
+    def _get_next_geometry(self, iteration: int) -> Tuple[str, Optional[str]]:
+        """Hook method to get geometry for next batch.
 
         Args:
             iteration: Current iteration number
+
+        Returns:
+            Tuple of (geom_path, optional_display_info)
+            - geom_path: Path to geometry file to use
+            - optional_display_info: Optional string to display (e.g., "Geometry: hex.obj")
         """
-        # Calculate minimum required orientations
-        min_required = self.min_batches * self.batch_size
+        # Default implementation: use fixed geometry from settings
+        return self.settings.geom_path, None
 
-        # Show progress with min orientations requirement
-        if self.n_orientations < min_required:
-            print(
-                f"\nIteration {iteration} ({self.n_orientations}/{min_required} orientations, min not reached):"
-            )
-        else:
-            print(
-                f"\nIteration {iteration} ({self.n_orientations} orientations, min {min_required} reached):"
-            )
+    def _handle_geometry_error(self, error: Exception, geom_path: str) -> bool:
+        """Hook method to handle geometry loading errors.
 
-        converged_status = self._check_convergence()
+        Args:
+            error: The exception that occurred
+            geom_path: Path to the geometry that failed
+
+        Returns:
+            True to skip this geometry and continue, False to raise the error
+        """
+        # Default implementation: re-raise error (fail fast for single geometry)
+        return False
+
+    def _get_theta_bins(self, variable: str) -> Optional[np.ndarray]:
+        """Get theta bins for Mueller elements from batch data."""
+        for batch in self.batch_data:
+            if "mueller_theta_bins" in batch:
+                return batch["mueller_theta_bins"]
 
+        # Fallback: infer from array length
+        mean_array, _ = self._calculate_mean_and_sem_array(variable)
+        if len(mean_array) > 0:
+            return np.linspace(0, 180, len(mean_array))
+
+        return None
+
+    def _update_convergence_history(self):
+        """Update convergence history with current SEM values."""
         for conv in self.convergables:
             if conv.is_mueller():
-                # Mueller element - show worst-case theta bin
+                # Mueller element - track worst SEM
                 mean_array, sem_array = self._calculate_mean_and_sem_array(
                     conv.variable
                 )
-
-                if len(mean_array) == 0:
-                    print(f"  {conv.variable:<10}: No data yet")
-                    continue
-
-                # Get theta bins from results (assuming we have access to bins_1d)
-                if hasattr(self, "settings") and hasattr(self.settings, "binning"):
-                    # We'll get theta values from the first batch's mueller_1d if available
-                    theta_bins = None
-                    for batch in self.batch_data:
-                        if "mueller_theta_bins" in batch:
-                            theta_bins = batch["mueller_theta_bins"]
-                            break
-                    if theta_bins is None:
-                        theta_bins = np.arange(len(mean_array))
-                else:
-                    theta_bins = np.arange(len(mean_array))
-
-                # Calculate relative SEM for each theta
-                if conv.tolerance_type == "relative":
-                    relative_sem_array = np.where(
-                        mean_array != 0, sem_array / np.abs(mean_array), float("inf")
-                    )
-                    worst_idx = np.argmax(relative_sem_array)
-                    worst_sem = relative_sem_array[worst_idx]
-                    target_str = f"{conv.tolerance * 100:.1f}%"
-                    current_str = f"{worst_sem * 100:.2f}%"
-                else:
-                    worst_idx = np.argmax(sem_array)
-                    worst_sem = sem_array[worst_idx]
-                    target_str = f"{conv.tolerance}"
-                    current_str = f"{worst_sem:.4g}"
-
-                worst_theta = theta_bins[worst_idx]
-                worst_mean = mean_array[worst_idx]
-
-                # Count converged bins (either all or specified indices)
-                if conv.theta_indices is not None:
-                    # Only checking specific bins
-                    indices = [i for i in conv.theta_indices if i < len(mean_array)]
+                if len(mean_array) > 0:
                     if conv.tolerance_type == "relative":
-                        converged_bins = np.sum(
-                            relative_sem_array[indices] < conv.tolerance
+                        relative_sem_array = np.where(
+                            mean_array != 0,
+                            sem_array / np.abs(mean_array),
+                            float("inf"),
                         )
+                        worst_sem = np.max(relative_sem_array)
                     else:
-                        converged_bins = np.sum(sem_array[indices] < conv.tolerance)
-                    total_bins = len(indices)
-                    bin_desc = (
-                        f"θ={[theta_bins[i] for i in indices]}"
-                        if len(indices) <= 3
-                        else f"{len(indices)} bins"
-                    )
-                else:
-                    # Checking all bins
-                    if conv.tolerance_type == "relative":
-                        converged_bins = np.sum(relative_sem_array < conv.tolerance)
-                    else:
-                        converged_bins = np.sum(sem_array < conv.tolerance)
-                    total_bins = len(mean_array)
-                    bin_desc = f"{total_bins} bins"
-
-                status = "✓" if converged_status[conv.variable] else "❌"
-
-                # Print Mueller convergence info
-                if conv.theta_indices is not None and len(conv.theta_indices) <= 3:
-                    # For small number of specific bins, show them explicitly
-                    print(
-                        f"  {conv.variable:<10}: {converged_bins}/{total_bins} {bin_desc} | "
-                        f"Worst θ={worst_theta:.1f}°: {worst_mean:.4g} | SEM: {current_str} (target: {target_str}) {status}"
-                    )
-                else:
-                    # For many bins, use standard format
-                    print(
-                        f"  {conv.variable:<10}: {converged_bins}/{total_bins} bins converged | "
-                        f"Worst θ={worst_theta:.1f}°: {worst_mean:.4g} | SEM: {current_str} (target: {target_str}) {status}"
-                    )
+                        worst_sem = np.max(sem_array)
 
-                # Add worst SEM to convergence history
-                self.convergence_history.append(
-                    (self.n_orientations, conv.variable, worst_sem)
-                )
+                    self.convergence_history.append(
+                        (self.n_orientations, conv.variable, worst_sem)
+                    )
             else:
                 # Scalar variable
                 mean, sem = self._calculate_mean_and_sem(conv.variable)
+                if conv.tolerance_type == "relative" and mean != 0:
+                    sem = sem / abs(mean)
 
-                # Calculate 95% CI
-                ci_lower = mean - 1.96 * sem
-                ci_upper = mean + 1.96 * sem
-
-                # Format based on tolerance type
-                if conv.tolerance_type == "relative":
-                    if mean != 0:
-                        relative_sem = sem / abs(mean)
-                        target_str = f"{conv.tolerance * 100:.1f}%"
-                        current_str = f"{relative_sem * 100:.2f}%"
-                    else:
-                        target_str = f"{conv.tolerance} (abs, mean=0)"
-                        current_str = f"{sem:.4g}"
-                else:
-                    target_str = f"{conv.tolerance}"
-                    current_str = f"{sem:.4g}"
-
-                # Status indicator
-                status = "✓" if converged_status[conv.variable] else "❌"
-
-                # Print line with mean, SEM, CI, and convergence status
-                print(
-                    f"  {conv.variable:<10}: {mean:.6f} ± {sem:.6f} [{ci_lower:.6f}, {ci_upper:.6f}] | SEM: {current_str} (target: {target_str}) {status}"
-                )
-
-                # Add to convergence history
                 self.convergence_history.append(
                     (self.n_orientations, conv.variable, sem)
                 )
@@ -577,30 +562,99 @@ class Convergence:
         converged = False
         warning = None
 
-        while not converged and self.n_orientations < self.max_orientations:
-            iteration += 1
-
-            # Determine batch size for this iteration
-            remaining = self.max_orientations - self.n_orientations
-            batch_size = min(self.batch_size, remaining)
-
-            # Set batch size
-            orientations = goad.create_uniform_orientation(batch_size)
-
-            # Set the orientations for the settings
-            self.settings.orientation = orientations
-
-            mp = goad.MultiProblem(self.settings)
-            mp.py_solve()
-
-            # Update statistics
-            self._update_statistics(mp.results, batch_size)
-
-            # Print progress
-            self._print_progress(iteration)
+        # Create Live context for smooth updating display
+        with self._display.create_live_context() as live:
+            # Show initial display before first batch
+            initial_display = self._display.build_display(
+                iteration=0,
+                n_orientations=self.n_orientations,
+                get_stats=self._calculate_mean_and_sem,
+                get_array_stats=self._calculate_mean_and_sem_array,
+                get_bin_labels=self._get_theta_bins,
+                power_ratio=None,
+                geom_info=None,
+            )
+            live.update(initial_display)
+
+            while not converged and self.n_orientations < self.max_orientations:
+                iteration += 1
+
+                # Get geometry for this batch (hook method - can be overridden)
+                geom_path, geom_info = self._get_next_geometry(iteration)
+
+                # Determine batch size for this iteration
+                remaining = self.max_orientations - self.n_orientations
+                batch_size = min(self.batch_size, remaining)
+
+                # Set batch size
+                orientations = goad.create_uniform_orientation(batch_size)
+
+                # Set the geometry path and orientations for the settings
+                self.settings.geom_path = geom_path
+                self.settings.orientation = orientations
+
+                # Run MultiProblem with error handling for bad geometries
+                # Suppress Rust progress bars by redirecting stderr at fd level
+                try:
+                    mp = goad.MultiProblem(self.settings)
+                    # Redirect stderr file descriptor to suppress Rust progress bars
+                    stderr_fd = sys.stderr.fileno()
+                    with open(os.devnull, "w") as devnull:
+                        old_stderr_fd = os.dup(stderr_fd)
+                        try:
+                            os.dup2(devnull.fileno(), stderr_fd)
+                            mp.py_solve()
+                        finally:
+                            os.dup2(old_stderr_fd, stderr_fd)
+                            os.close(old_stderr_fd)
+                except Exception as e:
+                    # Geometry loading failed (bad faces, degenerate geometry, etc.)
+                    # Check if subclass wants to handle this error (e.g., skip for ensemble)
+                    if self._handle_geometry_error(e, geom_path):
+                        # Skip this geometry and continue
+                        continue
+                    else:
+                        # For single-geometry convergence, we can't skip - must raise error
+                        error_msg = (
+                            f"Failed to initialize MultiProblem with geometry '{geom_path}': {e}\n"
+                            f"Please check geometry file for:\n"
+                            f"  - Degenerate faces (area = 0)\n"
+                            f"  - Non-planar geometry\n"
+                            f"  - Faces that are too small\n"
+                            f"  - Invalid mesh topology\n"
+                            f"  - Geometry file corruption"
+                        )
+                        raise type(e)(error_msg) from e
+
+                # Update statistics
+                self._update_statistics(mp.results, batch_size)
+
+                # Extract power ratio from results
+                try:
+                    powers = mp.results.powers  # It's a property, not a method
+                    power_in = powers.get("input", 1.0)
+                    power_out = powers.get("output", 0.0)
+                    power_ratio = power_out / power_in if power_in > 0 else 0.0
+                except Exception:
+                    power_ratio = None
+
+                # Update convergence history
+                self._update_convergence_history()
+
+                # Update live display with optional geometry info
+                display = self._display.build_display(
+                    iteration=iteration,
+                    n_orientations=self.n_orientations,
+                    get_stats=self._calculate_mean_and_sem,
+                    get_array_stats=self._calculate_mean_and_sem_array,
+                    get_bin_labels=self._get_theta_bins,
+                    power_ratio=power_ratio,
+                    geom_info=geom_info,
+                )
+                live.update(display)
 
-            # Check convergence
-            converged = self._all_converged()
+                # Check convergence
+                converged = self._all_converged()
 
         # Prepare final results
         if converged:
@@ -728,155 +782,49 @@ class EnsembleConvergence(Convergence):
             mueller_2d=mueller_2d,
         )
 
-    def run(self) -> ConvergenceResults:
-        """Run the ensemble convergence study.
+        # Track skipped geometries for error handling
+        self.skipped_geometries = []
 
-        Each batch iteration randomly selects a geometry file from the
-        ensemble directory before running the orientation averaging.
+    def _get_next_geometry(self, iteration: int) -> Tuple[str, Optional[str]]:
+        """Override to randomly select geometry from ensemble.
+
+        Args:
+            iteration: Current iteration number
 
         Returns:
-            ConvergenceResults containing final ensemble-averaged values
+            Tuple of (geom_path, display_info)
         """
-        iteration = 0
-        converged = False
-        warning = None
-
-        while not converged and self.n_orientations < self.max_orientations:
-            iteration += 1
-
-            # Randomly select a geometry file for this batch
-            geom_file = random.choice(self.geom_files)
-            geom_path = os.path.join(self.geom_dir, geom_file)
-
-            # Determine batch size for this iteration
-            remaining = self.max_orientations - self.n_orientations
-            batch_size = min(self.batch_size, remaining)
-
-            # Create orientations for this batch
-            orientations = goad.create_uniform_orientation(batch_size)
-
-            # Update settings with selected geometry and orientations
-            self.settings.geom_path = geom_path
-            self.settings.orientation = orientations
-
-            # Run MultiProblem with selected geometry
-            mp = goad.MultiProblem(self.settings)
-            mp.py_solve()
-
-            # Update statistics
-            self._update_statistics(mp.results, batch_size)
-
-            # Print progress (with geometry info)
-            min_required = self.min_batches * self.batch_size
-            if self.n_orientations < min_required:
-                print(
-                    f"\nIteration {iteration} ({self.n_orientations}/{min_required} orientations, min not reached) - Geometry: {geom_file}"
-                )
-            else:
-                print(
-                    f"\nIteration {iteration} ({self.n_orientations} orientations, min {min_required} reached) - Geometry: {geom_file}"
-                )
-            self._print_progress_without_header(iteration)
-
-            # Check convergence
-            converged = self._all_converged()
-
-        # Prepare final results
-        if converged:
-            print(f"\nConverged after {self.n_orientations} orientations.")
-        else:
-            warning = f"Maximum orientations ({self.max_orientations}) reached without convergence"
-            print(f"\nWarning: {warning}")
-
-        # Calculate final values and SEMs
-        final_values = {}
-        final_sems = {}
-        for conv in self.convergables:
-            if conv.is_mueller():
-                mean_array, sem_array = self._calculate_mean_and_sem_array(
-                    conv.variable
-                )
-                final_values[conv.variable] = mean_array
-                final_sems[conv.variable] = sem_array
-            else:
-                mean, sem = self._calculate_mean_and_sem(conv.variable)
-                final_values[conv.variable] = mean
-                final_sems[conv.variable] = sem
-
-        # Prepare Mueller matrices with SEM
-        mueller_1d = None
-        mueller_1d_sem = None
-        mueller_2d = None
-
-        if self.mueller_1d and self.mueller_1d_sum is not None:
-            mueller_1d = self.mueller_1d_sum / self.n_orientations
-
-            # Compute SEM for all 16 Mueller elements
-            # mueller_1d shape: (n_theta, 16)
-            n_theta = mueller_1d.shape[0]
-            mueller_1d_sem = np.zeros_like(mueller_1d)
-
-            for row in range(1, 5):
-                for col in range(1, 5):
-                    element_name = f"S{row}{col}"
-                    mueller_idx = (row - 1) * 4 + (col - 1)
-
-                    # Calculate mean and SEM for this element across all theta bins
-                    mean_array, sem_array = self._calculate_mean_and_sem_array(
-                        element_name
-                    )
-
-                    if len(sem_array) > 0:
-                        mueller_1d_sem[:, mueller_idx] = sem_array
-
-            # Store mueller_1d_sem in final_values for unified API access
-            final_values["mueller_1d_sem"] = mueller_1d_sem
+        # Randomly select a geometry file for this batch
+        geom_file = random.choice(self.geom_files)
+        geom_path = os.path.join(self.geom_dir, geom_file)
+        geom_info = f"Geom: {geom_file}"
 
-        if self.mueller_2d and self.mueller_2d_sum is not None:
-            mueller_2d = self.mueller_2d_sum / self.n_orientations
+        return geom_path, geom_info
 
-        return ConvergenceResults(
-            converged=converged,
-            n_orientations=self.n_orientations,
-            values=final_values,
-            sem_values=final_sems,
-            mueller_1d=mueller_1d,
-            mueller_2d=mueller_2d,
-            convergence_history=self.convergence_history,
-            warning=warning,
-        )
+    def _handle_geometry_error(self, error: Exception, geom_path: str) -> bool:
+        """Override to skip bad geometries in ensemble mode.
 
-    def _print_progress_without_header(self, iteration: int):
-        """Print convergence progress without iteration header (already printed with geometry)."""
-        converged_status = self._check_convergence()
+        Args:
+            error: The exception that occurred
+            geom_path: Path to the geometry that failed
 
-        for conv in self.convergables:
-            mean, sem = self._calculate_mean_and_sem(conv.variable)
-
-            # Calculate 95% CI
-            ci_lower = mean - 1.96 * sem
-            ci_upper = mean + 1.96 * sem
-
-            # Format based on tolerance type
-            if conv.tolerance_type == "relative":
-                if mean != 0:
-                    relative_sem = sem / abs(mean)
-                    target_str = f"{conv.tolerance * 100:.1f}%"
-                    current_str = f"{relative_sem * 100:.2f}%"
-                else:
-                    target_str = f"{conv.tolerance} (abs, mean=0)"
-                    current_str = f"{sem:.4g}"
-            else:
-                target_str = f"{conv.tolerance}"
-                current_str = f"{sem:.4g}"
+        Returns:
+            True to skip this geometry and continue
+        """
+        # Extract just the filename
+        geom_file = os.path.basename(geom_path)
 
-            # Status indicator
-            status = "✓" if converged_status[conv.variable] else "❌"
+        # Print warning and track skipped geometry
+        print(f"\nWarning: Skipping geometry '{geom_file}': {error}")
+        self.skipped_geometries.append(geom_file)
 
-            # Print line with mean, SEM, CI, and convergence status
-            print(
-                f"  {conv.variable:<10}: {mean:.6f} ± {sem:.6f} [{ci_lower:.6f}, {ci_upper:.6f}] | SEM: {current_str} (target: {target_str}) {status}"
+        # Check if all geometries have been skipped
+        if len(self.skipped_geometries) >= len(self.geom_files):
+            raise ValueError(
+                f"All {len(self.geom_files)} geometry files failed to load. "
+                "Please check geometry files for degenerate faces, non-planar geometry, "
+                "or faces that are too small."
             )
 
-            # Add to convergence history
-            self.convergence_history.append((self.n_orientations, conv.variable, sem))
+        # Skip this geometry and continue
+        return True