Feat/extended sources (#470)

* feat: implement extended source modelling logic and visualization (#224)

* feat: improve source visualization in optical system and source.draw() method (#224)

* correct formulas for creating the collimated gaussian source, and add a notebook as an example - to me moved to the gallery (#224)

* formatting

* refactor source logic: Add SMF Source model, ExtendedSourceOptic wrapper, and a correction term added for radiant intensity computation (#224)

* ruff check

* docs: update documentation to include ExtendedSource approach. Includes an example Notebook Tutorial and for the gallery as well

* Add comprehensive tests for extended sources and fix torch XY projection

- Add tests/test_extended_sources.py with 160 tests covering:
  - SMFSource: ray generation, parameter computation, edge cases, visualization
  - BaseSource: abstract interface contract verification
  - ExtendedSourceOptic: delegation, tracing, drawing, projections, repr
  - IncoherentIrradiance: source integration and conflict handling
  - SourceViewer: 6-panel diagnostic visualization
  - Standard analysis delegation via .optic for field-based analyses

- Fix torch backend compatibility in Lens2D XY projection
  (optiland/visualization/system/lens.py): convert tensor values to
  plain floats via be.to_numpy().item() before passing to plt.Circle,
  preventing RuntimeError from matplotlib calling .numpy() on tensors
  with requires_grad=True

* chore: trying to resolve doc gen failure. Remove autogen'd files

* chore: alphabetize api

* chore: run notebook

* docs: add extended sources to MD learning guide

---------

Co-authored-by: Kramer <kdanielharrison@gmail.com>

Author: manuelFragata

docs/LEARNING_GUIDE.md

@@ -99,14 +99,17 @@ Welcome to the Optiland Learning Guide! This guide walks you through key concept
         - Adding new coating types
     - [Tutorial 10c - Custom Optimization Algorithms](https://github.com/HarrisonKramer/optiland/blob/master/docs/examples/Tutorial_10c_Custom_Optimization_Algorithm.ipynb)
         - Creating a "random walk optimizer" to optimize an aspheric singlet
-11. **Machine Learning in Optical Design** - note that these notebooks are hosted in the [LensAI repository](https://github.com/HarrisonKramer/LensAI)
-    - [Tutorial 11a - Random Forest Regressor to Predict Optimal Lens Properties](https://github.com/HarrisonKramer/LensAI/blob/main/notebooks/Example_1/Singlet_RF_Model_RMS_Spot_Size.ipynb)
+11. **Extended Sources**
+    - [Tutorial 11a - Extended Source Modeling](https://github.com/HarrisonKramer/optiland/blob/master/docs/examples/Tutorial_11a_Extended_Source_Modeling.ipynb)
+        - Modeling extended sources
+12. **Machine Learning in Optical Design** - note that these notebooks are hosted in the [LensAI repository](https://github.com/HarrisonKramer/LensAI)
+    - [Tutorial 12a - Random Forest Regressor to Predict Optimal Lens Properties](https://github.com/HarrisonKramer/LensAI/blob/main/notebooks/Example_1/Singlet_RF_Model_RMS_Spot_Size.ipynb)
         - Demonstrates how to build and train a random forest regressor to predict the radius of curvature of a plano-convex lens in order to minimize the RMS spot size.
-    - [Tutorial 11b - Ray Path Failure Classification Model](https://github.com/HarrisonKramer/LensAI/blob/main/notebooks/Example_2/Ray_Path_Failure_Classification_Model.ipynb)
+    - [Tutorial 12b - Ray Path Failure Classification Model](https://github.com/HarrisonKramer/LensAI/blob/main/notebooks/Example_2/Ray_Path_Failure_Classification_Model.ipynb)
         - Uses logistic regression to predict ray path failures in a Cooke triplet design.
-    - [Tutorial 11c - Surrogate Ray Tracing Model Using Neural Networks](https://github.com/HarrisonKramer/LensAI/blob/main/notebooks/Example_3/Double_Gauss_Surrogate_Model.ipynb)
+    - [Tutorial 12c - Surrogate Ray Tracing Model Using Neural Networks](https://github.com/HarrisonKramer/LensAI/blob/main/notebooks/Example_3/Double_Gauss_Surrogate_Model.ipynb)
         - Builds a neural network surrogate ray tracing model to increase effective "ray tracing" speed by 10,000x.
-    - [Tutorial 11d - Super-Resolution Generative Adversarial Network to Enhance Wavefront Map Data](https://github.com/HarrisonKramer/LensAI/blob/main/notebooks/Example_5/SR_GAN_for_wavefront_data.ipynb)
+    - [Tutorial 12d - Super-Resolution Generative Adversarial Network to Enhance Wavefront Map Data](https://github.com/HarrisonKramer/LensAI/blob/main/notebooks/Example_5/SR_GAN_for_wavefront_data.ipynb)
         - Utilizes a super-resolution GAN (SRGAN) to upscale low-resolution wavefront data into high-resolution data.
-    - [Tutorial 11e - Optimization of Aspheric Lenses via Reinforcement Learning](https://github.com/HarrisonKramer/LensAI/blob/main/notebooks/Example_4/RL_aspheric_singlet.ipynb)
+    - [Tutorial 12e - Optimization of Aspheric Lenses via Reinforcement Learning](https://github.com/HarrisonKramer/LensAI/blob/main/notebooks/Example_4/RL_aspheric_singlet.ipynb)
         - Reinforcement learning is applied to the optimization of aspheric singlet lenses to generate new lens designs.

docs/api/api_optic.rst

@@ -3,10 +3,12 @@ Optic
 
 This section includes the core optic module and the optic updater module. The Optic class is the primary container for an optical system in Optiland.
 The OpticUpdater class is used to update various properties of an Optic object, such as surface properties or polarization.
+The ExtendedSourceOptic class is a wrapper around the Optic class that enables extended source ray tracing.
 
 .. autosummary::
    :toctree: optic/
    :caption: Optic Modules
 
    optic.optic
    optic.optic_updater
+   optic.extended_source_optic

docs/api/api_sources.rst

@@ -0,0 +1,19 @@
+Sources
+=======
+
+This section covers the extended source modeling modules of the Optiland package.
+Extended sources allow users to define spatially and angularly extended light sources
+(such as single-mode fiber outputs) and trace them through optical systems.
+
+The :class:`sources.base.BaseSource` class defines the abstract interface that all
+source implementations must follow. Concrete implementations, such as
+:class:`sources.smf.SMFSource`, provide specific source types with their own
+spatial and angular distributions.
+
+.. autosummary::
+   :toctree: sources/
+   :caption: Sources Modules
+
+   sources.base
+   sources.smf
+   sources.visualization

docs/examples/Tutorial_11a_Extended_Source_Modeling.ipynb

@@ -49,6 +49,20 @@ Optimization and Tolerancing
 - **Extensible Framework**:
   Add new optimization variables, constraints, or algorithms with minimal overhead.
 
+Extended Source Modeling
+-----------------------
+
+- **Extended Source Ray Tracing**:
+  Model spatially and angularly extended light sources and trace them through optical systems using the ``ExtendedSourceOptic`` wrapper.
+- **Single-Mode Fiber (SMF) Source**:
+  Generate physically accurate ray bundles representing fiber outputs, with Gaussian spatial and angular distributions and quasi-random Sobol sampling.
+- **Custom Source Support**:
+  Define custom source types by inheriting from ``BaseSource`` and implementing the ``generate_rays`` method.
+- **Source Visualization**:
+  Validate source definitions with multi-panel diagnostic plots showing spatial distributions, angular distributions, cross-sections, and ray propagation paths.
+- **Irradiance Analysis Compatibility**:
+  Compute incoherent irradiance distributions at detector surfaces from extended source illumination.
+
 Material Database
 -----------------
 

docs/functionalities.rst

@@ -0,0 +1,12 @@
+.. _gallery_extended_sources:
+
+Extended Sources
+================
+
+This section contains examples of extended source modeling in **Optiland**. Extended sources allow
+users to define spatially and angularly extended light sources and trace them through optical systems,
+enabling physically accurate simulations of fiber-coupled and other non-point-source illumination.
+
+.. nbgallery::
+
+    extended_sources/beam_shaping_singlet

docs/gallery/extended_sources.rst

@@ -108,6 +108,7 @@ Not sure what to type in the shell? Here are a few ideas to explore Optiland rig
    gallery/differentiable_ray_tracing
    gallery/real_world_projects
    gallery/external_tools
+   gallery/extended_sources
    gallery/miscellaneous
 
 
@@ -175,6 +176,7 @@ Not sure what to type in the shell? Here are a few ideas to explore Optiland rig
    api/api_rays
    api/api_raytrace
    api/api_solves
+   api/api_sources
    api/api_surfaces
    api/api_tolerancing
    api/api_visualization

docs/gallery/extended_sources/beam_shaping_singlet.ipynb

@@ -114,7 +114,15 @@ This learning guide breaks down Optiland into a series of tutorials that cover t
    examples/Tutorial_10b_Custom_Coating_Types
    examples/Tutorial_10c_Custom_Optimization_Algorithm
 
-11. Machine Learning in Optical Design
+11. Extended Source Modeling
+---------------------------
+
+.. toctree::
+   :maxdepth: 1
+
+   examples/Tutorial_11a_Extended_Source_Modeling
+
+12. Machine Learning in Optical Design
 --------------------------------------
 
 These examples demonstrate how Optiland can be used in conjunction with machine learning to solve optical design problems.

docs/index.rst

@@ -48,6 +48,10 @@ class RadiantIntensity(BaseAnalysis):
         num_rays (int): Number of rays to trace if user_initial_rays is None.
         distribution_name (str): Ray distribution if user_initial_rays is None.
         user_initial_rays (RealRays | None): Optional user-provided initial rays.
+        source (BaseSource | None): Optional extended source object
+            (e.g., GaussianSource) to generate initial rays automatically.
+            Cannot be used with user_initial_rays. When provided, num_rays
+            determines how many rays to generate.
         data (list[list[tuple]]): Stores (intensity_map,
                                           angle_X_bin_edges, angle_Y_bin_edges,
                                           angle_X_bin_centers, angle_Y_bin_centers)
@@ -73,6 +77,8 @@ def __init__(
         num_rays: int = 100000,
         distribution: DistributionType = "random",
         user_initial_rays=None,
+        source=None,
+        skip_trace: bool = False,
     ):
         if fields == "all":
             self.fields = optic.fields.get_field_coords()
@@ -81,6 +87,31 @@ def __init__(
                 fields = [fields]
             self.fields = tuple(fields)
 
+        # Handle source integration
+        if source is not None and user_initial_rays is not None:
+            raise ValueError("Cannot specify both 'source' and 'user_initial_rays'.")
+
+        self.user_initial_rays = user_initial_rays
+        if source is not None:
+            # Generate rays from the extended source
+            self.user_initial_rays = source.generate_rays(num_rays)
+            # When using a source, we treat all rays as a single "field"
+            # The source emission defines the field, not optic.fields
+            self.fields = [(0.0, 0.0)]  # Single dummy field for source rays
+
+        self._initial_ray_data = None
+        if self.user_initial_rays is not None:
+            self._initial_ray_data = {
+                "x": self.user_initial_rays.x,
+                "y": self.user_initial_rays.y,
+                "z": self.user_initial_rays.z,
+                "L": self.user_initial_rays.L,
+                "M": self.user_initial_rays.M,
+                "N": self.user_initial_rays.N,
+                "intensity": self.user_initial_rays.i,
+                "wavelength": self.user_initial_rays.w,
+            }
+
         self.num_angular_bins_X = num_angular_bins_X
         self.num_angular_bins_Y = num_angular_bins_Y
         self.angle_X_min, self.angle_X_max = float(angle_X_min), float(angle_X_max)
@@ -89,7 +120,7 @@ def __init__(
         # for absolute units, we need to ensure the user has provided rays
         # with 'calibrated' power
         self.use_absolute_units = use_absolute_units
-        if self.use_absolute_units and user_initial_rays is None:
+        if self.use_absolute_units and self.user_initial_rays is None:
             print(
                 "Warning: `use_absolute_units` is True, but no `user_initial_rays` "
                 "were provided."
@@ -103,7 +134,7 @@ def __init__(
         self.reference_surface_index = int(reference_surface_index)
         self.num_rays = num_rays
         self.distribution_name: DistributionType = distribution
-        self.user_initial_rays = user_initial_rays
+        self.skip_trace = skip_trace
 
         super().__init__(optic, wavelengths)
 
@@ -121,15 +152,19 @@ def _generate_data(self):
     def _generate_field_wavelength_data(
         self, field_coord: tuple[ScalarOrArray, ScalarOrArray], wavelength: float
     ) -> tuple[BEArray, BEArray, BEArray, BEArray, BEArray]:
-        if self.user_initial_rays is None:
-            self.optic.trace(
-                *field_coord,
-                wavelength=wavelength,
-                num_rays=self.num_rays,
-                distribution=self.distribution_name,
-            )
-        else:
-            self.optic.surface_group.trace(self.user_initial_rays)
+        if not self.skip_trace:
+            if self.user_initial_rays is None:
+                self.optic.trace(
+                    *field_coord,
+                    wavelength=wavelength,
+                    num_rays=self.num_rays,
+                    distribution=self.distribution_name,
+                )
+            else:
+                from optiland.rays import RealRays
+
+                rays_to_trace = RealRays(**self._initial_ray_data)
+                self.optic.surface_group.trace(rays_to_trace)
 
         surf_group = self.optic.surface_group
         try:
@@ -200,13 +235,43 @@ def _generate_field_wavelength_data(
                 )
 
         if self.use_absolute_units:
-            delta_angle_X_rad = be.radians(angle_X_bins[1] - angle_X_bins[0])
-            delta_angle_Y_rad = be.radians(angle_Y_bins[1] - angle_Y_bins[0])
-            solid_angle_per_bin_sr = delta_angle_X_rad * delta_angle_Y_rad
-
+            # 1. Calculate basic bin sizes (radians)
+            dx = be.radians(angle_X_bins[1] - angle_X_bins[0])
+            dy = be.radians(angle_Y_bins[1] - angle_Y_bins[0])
+
+            # 2. Create meshgrid of bin centers (in Radians) for
+            # the Jacobian calculation
+            # Note: We must be careful with tensor shapes here to match
+            # power_map (Y, X)
+            ax_c_rad = be.radians(angle_X_centers)
+            ay_c_rad = be.radians(angle_Y_centers)
+
+            # Create grids. Meshgrid usually returns (Y, X) with
+            # indexing='ij' or 'xy' depending on backend
+            # For safety, let's explicitely broadcast
+            # resulting shape (Y, X) usually
+            AX, AY = be.meshgrid(ax_c_rad, ay_c_rad)
+
+            # 3. Compute Jacobian terms
+            # J = (sec^2(tx) * sec^2(ty)) / (1 + tan^2(tx) + tan^2(ty))^(3/2)
+            tan2_tx = be.tan(AX) ** 2
+            tan2_ty = be.tan(AY) ** 2
+            sec2_tx = 1.0 + tan2_tx  # Identity: sec^2 = 1 + tan^2
+            sec2_ty = 1.0 + tan2_ty
+
+            numerator = sec2_tx * sec2_ty
+            denominator = (1.0 + tan2_tx + tan2_ty) ** 1.5
+
+            jacobian_factor = numerator / denominator
+
+            # 4. Compute true solid angle per bin
+            # d_omega = J * d_theta_x * d_theta_y
+            true_solid_angle_map = jacobian_factor * dx * dy
+
+            # 5. Normalize Power Map
             final_intensity_map = be.where(
-                solid_angle_per_bin_sr > 1e-12,
-                power_map / solid_angle_per_bin_sr,
+                true_solid_angle_map > 1e-12,
+                power_map / true_solid_angle_map,
                 be.zeros_like(power_map),
             )
         else: