Feat: Add PointDipole.from_angles() (#2623)

* Feat: Add `PointDipole.from_angles()`

Class method that returns a list of `PointDipole` objects used to emulate a single dipole polarized in an arbitrary direction. The direction is specificed using a polar and azimuthal angle.

* Fix: Pass polarization argument as string + NumPy-style docstring.

* Test: Simple testing function for `PointDipole.from_angles()`.

* Refc: Rename `PointDipole.from_angles()` to `PointDipole.sources_from_angles()`

---------

Co-authored-by: momchil-flex <[email protected]>

Author: saeedece

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Add support for `np.unwrap` in `tidy3d.plugins.autograd`.
 - Add Nunley variant to germanium material library based on Nunley et al. 2016 data.
+- Add `PointDipole.sources_from_angles()` that constructs a list of `PointDipole` objects needed to emulate a dipole oriented at a user-provided set of polar and azimuthal angles.
 - Added `priority` parameter to `web.run()` and related functions to allow vGPU users to set task priority (1-10) in the queue.
 
 ### Changed

tests/test_components/test_source.py

@@ -148,6 +148,67 @@ def test_dipole():
         _ = td.PointDipole(size=(1, 1, 1), source_time=g, center=(1, 2, 3), polarization="Ex")
 
 
+def test_dipole_sources_from_angles():
+    g = td.GaussianPulse(freq0=1e12, fwidth=0.1e12)
+
+    with pytest.raises(pydantic.ValidationError):
+        _ = td.PointDipole.sources_from_angles(
+            size=(1, 1, 1),
+            source_time=g,
+            center=(1, 2, 3),
+            angle_theta=np.pi / 4,
+            angle_phi=np.pi / 4,
+        )
+
+    with pytest.raises(ValueError):
+        _ = td.PointDipole.sources_from_angles(
+            source_time=g,
+            angle_theta=np.pi / 4,
+            angle_phi=np.pi / 4,
+            component="invalid",
+            center=(1, 2, 3),
+        )
+
+    assert (
+        len(
+            td.PointDipole.sources_from_angles(
+                source_time=g,
+                angle_theta=np.pi / 4,
+                angle_phi=np.pi / 4,
+                component="electric",
+                center=(1, 2, 3),
+            )
+        )
+        == 3
+    )
+
+    assert (
+        len(
+            td.PointDipole.sources_from_angles(
+                source_time=g,
+                angle_theta=np.pi / 4,
+                angle_phi=np.pi / 2,
+                component="electric",
+                center=(1, 2, 3),
+            )
+        )
+        == 2
+    )
+
+    assert (
+        len(
+            td.PointDipole.sources_from_angles(
+                source_time=g,
+                angle_theta=np.pi / 2,
+                angle_phi=np.pi / 2,
+                component="electric",
+                center=(1, 2, 3),
+            )
+        )
+        == 1
+    )
+
+
 def test_FieldSource():
     g = td.GaussianPulse(freq0=1e12, fwidth=0.1e12)
     mode_spec = td.ModeSpec(num_modes=2)

tidy3d/components/source/current.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from abc import ABC
+from math import cos, isclose, sin
 from typing import Optional
 
 import pydantic.v1 as pydantic
@@ -16,6 +17,7 @@
 from tidy3d.constants import MICROMETER
 
 from .base import Source
+from .time import SourceTimeType
 
 
 class CurrentSource(Source, ABC):
@@ -107,6 +109,62 @@ class PointDipole(CurrentSource, ReverseInterpolatedSource):
         units=MICROMETER,
     )
 
+    @classmethod
+    def sources_from_angles(
+        cls,
+        source_time: SourceTimeType,
+        angle_theta: float,
+        angle_phi: float,
+        component: Literal["electric", "magnetic"] = "electric",
+        **kwargs,
+    ) -> list[PointDipole]:
+        """Returns a list of `PointDipole` objects used to emulate a single dipole polarized in an arbitrary direction. The direction is specificed using a polar and azimuthal angle.
+
+        Parameters
+        ----------
+        source_time: :class:`.SourceTime`
+            Specification of the source time-dependence.
+        angle_theta : float
+            Polar angle w.r.t. the z-axis.
+        angle_phi : float
+            Azimuth angle around the z-axis.
+        component : Literal["electric", "magnetic"] = "electric"
+            The type of polarization.
+        kwargs : dict
+            Keyword arguments passed to ``PointDipole()``, excluding ``source_time`` and ``polarization``
+
+        Returns
+        -------
+        list[PointDipole]
+            A list of ``PointDipole`` objects that emulate a single dipole with an arbitrary direction of polarization.
+        """
+        if not (component == "electric" or component == "magnetic"):
+            raise ValueError('Component must be "electric" or "magnetic"')
+
+        dipoles: list[PointDipole] = []
+        polarizations = ["Ex", "Ey", "Ez"] if component == "electric" else ["Hx", "Hy", "Hz"]
+
+        multipliers = [
+            sin(angle_theta) * cos(angle_phi),
+            sin(angle_theta) * sin(angle_phi),
+            cos(angle_theta),
+        ]
+
+        for polarization, mult in zip(polarizations, multipliers):
+            if not isclose(mult, 0.0, rel_tol=0.0, abs_tol=1e-9):
+                modulated_source_time = source_time.updated_copy(
+                    amplitude=source_time.amplitude * mult
+                )
+                dipoles.append(
+                    cls(
+                        source_time=modulated_source_time,
+                        polarization=polarization,
+                        **kwargs,
+                    )
+                )
+
+        return dipoles
+
 
 class CustomCurrentSource(ReverseInterpolatedSource):
     """Implements a source corresponding to an input dataset containing ``E`` and ``H`` fields.