Allow setting radius in Helioprojective.assume_spherical_screen

sunpy/coordinates/frames.py

@@ -679,14 +679,14 @@ def is_visible(self, *, tolerance: u.m = 1*u.m):
 
     @classmethod
     @contextmanager
-    def assume_spherical_screen(cls, center, only_off_disk=False):
+    def assume_spherical_screen(cls, center, only_off_disk=False, radius='center_to_sun'):
         """
         Context manager to interpret 2D coordinates as being on the inside of a spherical screen.
 
-        The radius of the screen is the distance between the specified ``center`` and Sun center.
-        This ``center`` does not have to be the same as the observer location for the coordinate
-        frame.  If they are the same, then this context manager is equivalent to assuming that the
-        helioprojective "zeta" component is zero.
+        The radius of the screen is, by default, the distance between the specified ``center`` and
+        Sun center. This ``center`` does not have to be the same as the observer location for the
+        coordinate frame.  If they are the same, then this context manager is equivalent to assuming
+        that the helioprojective "zeta" component is zero.
 
         This replaces the default assumption where 2D coordinates are mapped onto the surface of the
         Sun.
@@ -698,6 +698,9 @@ def assume_spherical_screen(cls, center, only_off_disk=False):
         only_off_disk : `bool`, optional
             If `True`, apply this assumption only to off-disk coordinates, with on-disk coordinates
             still mapped onto the surface of the Sun.  Defaults to `False`.
+        radius : `~astropy.units.Quantity`
+            The radius of the spherical screen. The default sets the radius to the distance from the
+            screen center to the Sun.
 
         Examples
         --------
@@ -733,10 +736,13 @@ def assume_spherical_screen(cls, center, only_off_disk=False):
         try:
             old_spherical_screen = cls._spherical_screen  # nominally None
 
-            center_hgs = center.transform_to(HeliographicStonyhurst(obstime=center.obstime))
+            if radius == 'center_to_sun':
+                center_hgs = center.transform_to(HeliographicStonyhurst(obstime=center.obstime))
+                radius = center_hgs.radius
+
             cls._spherical_screen = {
                 'center': center,
-                'radius': center_hgs.radius,
+                'radius': radius,
                 'only_off_disk': only_off_disk
             }
             yield

sunpy/coordinates/tests/test_transformations.py

@@ -120,6 +120,45 @@ def test_hpc_hpc_null():
     assert hpc_out.observer == hpc_new.observer
 
 
+def test_hpc_hpc_spherical_screen():
+    D0 = 20*u.R_sun
+    L0 = 67.5*u.deg
+    Tx0 = 45*u.deg
+    observer_in = HeliographicStonyhurst(lat=0*u.deg, lon=0*u.deg, radius=D0)
+    # Once our coordinate is placed on the screen, observer_out will be looking
+    # directly along the line containing itself, the coordinate, and the Sun
+    observer_out = HeliographicStonyhurst(lat=0*u.deg, lon=L0, radius=2*D0)
+
+    sc_in = SkyCoord(Tx0, 0*u.deg, observer=observer_in,
+                     frame='helioprojective')
+
+    with Helioprojective.assume_spherical_screen(observer_in):
+        sc_3d = sc_in.make_3d()
+        sc_out = sc_in.transform_to(Helioprojective(observer=observer_out))
+
+    assert quantity_allclose(sc_3d.distance, D0)
+
+    assert quantity_allclose(sc_out.Tx, 0*u.deg, atol=1e-6*u.deg)
+    assert quantity_allclose(sc_out.Ty, 0*u.deg, atol=1e-6*u.deg)
+    # Law of Cosines to compute the coordinate's distance from the Sun, and then
+    # r_expected is the distance from observer_out to the coordinate
+    radius_expected = 2 * D0 - np.sqrt(2 * D0**2 - 2 * D0 * D0 * np.cos(45*u.deg))
+    assert quantity_allclose(sc_out.distance, radius_expected)
+
+    # Now test with a very large screen, letting us approximate the two
+    # observers as being the same (aside from a different zero point for Tx)
+    r_s = 1e9 * u.lightyear
+    with Helioprojective.assume_spherical_screen(observer_in, radius=r_s):
+        sc_3d = sc_in.make_3d()
+        sc_out = sc_in.transform_to(Helioprojective(observer=observer_out))
+
+    assert quantity_allclose(sc_3d.distance, r_s)
+
+    assert quantity_allclose(sc_out.Tx, Tx0 + L0)
+    assert quantity_allclose(sc_out.Ty, 0*u.deg)
+    assert quantity_allclose(sc_out.distance, r_s)
+
+
 def test_hcrs_hgs():
     # Get the current Earth location in HCRS
     adate = parse_time('2015/05/01 01:13:00')