From 39ef3af656d2e194e07a2b286dbd08342315efb7 Mon Sep 17 00:00:00 2001
From: anand43 <anand43@purdue.edu>
Date: Wed, 15 Nov 2023 14:27:26 -0500
Subject: [PATCH] edited to give the user to option to set a refernce radius
 for the CB

---
 swiftest/sph_harmonics.py | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/swiftest/sph_harmonics.py b/swiftest/sph_harmonics.py
index 6daa4f93c..25f74202a 100644
--- a/swiftest/sph_harmonics.py
+++ b/swiftest/sph_harmonics.py
@@ -29,7 +29,7 @@
 )
 
 class Sph_Harmonics(object):
-    def clm_from_ellipsoid(mass, density, a, b = None, c = None, lmax = 6, ref_radius = False):
+    def clm_from_ellipsoid(mass, density, a, b = None, c = None, lmax = 6, lref_radius = False, ref_radius = None):
         """
         Creates and returns the gravity coefficients for an ellipsoid with principal axes a, b, c upto a certain maximum degree lmax. 
         Uses pyshtools. No units necessary for a, b, & c. However, they need to be in the same units (DU).
@@ -48,8 +48,10 @@ def clm_from_ellipsoid(mass, density, a, b = None, c = None, lmax = 6, ref_radiu
             length of the pricipal axis aligned with the z axis
         lmax : int, optional, default = 6
             The maximum spherical harmonic degree resolvable by the grid.
-        ref_radius : boolean, optional, default = False
+        lref_radius : boolean, optional, default = False
             Boolean value to return the reference radius calculated by SHTOOLS
+        ref_radius : float, optional, default = None
+            Reference radius to scale the gravitational coefficients to
 
         Returns
         -------
@@ -77,13 +79,16 @@ def clm_from_ellipsoid(mass, density, a, b = None, c = None, lmax = 6, ref_radiu
         # Return reference radius EQUALS the radius of the Central Body
         print(f'Ensure that the Central Body radius equals the reference radius.')
 
-        if(ref_radius == True):
+        if(lref_radius == True and ref_radius is None):
             ref_radius = shape_SH.expand(normalization = '4pi').coeffs[0, 0, 0]
             return clm, ref_radius
+        else if(lref_radius == True and ref_radius is not None):
+            clm = clm.change_ref(r0 = ref_radius)
+            return clm, ref_radius
         else:
             return clm
 
-    def clm_from_relief(mass, density, grid, lmax = 6, ref_radius = True):
+    def clm_from_relief(mass, density, grid, lmax = 6, lref_radius = False, ref_radius = None):
         """
         Creates and returns the gravity coefficients for a body with a given DH grid upto a certain maximum degree lmax. 
         Uses pyshtools.
@@ -98,8 +103,10 @@ def clm_from_relief(mass, density, grid, lmax = 6, ref_radius = True):
             DH grid of the surface relief of the body
         lmax : int, optional, default = 6
             The maximum spherical harmonic degree resolvable by the grid.
-        ref_radius : boolean, optional, default = False
+        lref_radius : boolean, optional, default = False
             Boolean value to return the reference radius calculated by SHTOOLS
+        ref_radius : float, optional, default = None
+            Reference radius to scale the gravitational coefficients to
 
         Returns
         -------
@@ -129,8 +136,11 @@ def clm_from_relief(mass, density, grid, lmax = 6, ref_radius = True):
 
         print(f'Ensure that the Central Body radius equals the reference radius.')
 
-        if(ref_radius == True):
+        if(lref_radius == True and ref_radius is None):
             ref_radius = shape_SH.expand(normalization = '4pi').coeffs[0, 0, 0]
             return clm, ref_radius
+        else if(lref_radius == True and ref_radius is not None):
+            clm = clm.change_ref(r0 = ref_radius)
+            return clm, ref_radius
         else:
             return clm