From 7cfbff2265af6d7667ca37598ec87dedc18154ef Mon Sep 17 00:00:00 2001
From: jgieseler <jan.gieseler@utu.fi>
Date: Thu, 21 Sep 2023 17:41:29 +0300
Subject: [PATCH] extend pfss_3d options, add reference_long to it

---
 solarmach/__init__.py | 73 +++++++++++++++++++++++++++++++++++++++----
 1 file changed, 67 insertions(+), 6 deletions(-)

diff --git a/solarmach/__init__.py b/solarmach/__init__.py
index 1de8677..6057349 100644
--- a/solarmach/__init__.py
+++ b/solarmach/__init__.py
@@ -1376,7 +1376,8 @@ def legend_arrow(width, height, **_):
             return fig, ax
 
     def pfss_3d(self, active_area=(None, None, None, None), color_code='object', rss=2.5,
-                plot_spirals=True, plot_sun_body_line=False, numbered_markers=False, zoom_out=False):
+                plot_spirals=True, plot_sun_body_line=False, numbered_markers=False, plot_equatorial_plane=True,
+                reference_vsw=400, zoom_out=False):
         """
         https://plotly.github.io/plotly.py-docs/generated/plotly.graph_objects.Scatter3d.html
         https://plotly.com/python-api-reference/generated/plotly.graph_objects.Figure.html
@@ -1681,11 +1682,71 @@ def pfss_3d(self, active_area=(None, None, None, None), color_code='object', rss
                               margin=dict(r=20, l=10, b=10, t=10),
                               )
 
-        fig.add_trace(go.Surface(x=np.linspace(-200, 200, 100),
-                                 y=np.linspace(-200, 200, 100),
-                                 z=np.zeros((100, 100)),
-                                 hoverinfo='skip',
-                                 colorscale='gray', showscale=False, opacity=0.2))
+        if self.reference_long is not None:
+            delta_ref = self.reference_long
+            if delta_ref < 0.:
+                delta_ref = delta_ref + 360.
+            if self.reference_lat is None:
+                ref_lat = 0.
+            else:
+                ref_lat = self.reference_lat
+
+            omega_ref = self.solar_diff_rot(ref_lat)
+            alpha_ref = np.deg2rad(delta_ref) + omega_ref / (reference_vsw / AU) * (self.target_solar_radius*aconst.R_sun.to(u.AU).value - r_array) * np.cos(np.deg2rad(ref_lat))
+
+            arrow_dist = min([self.max_dist/3.2, 2.])
+            x, y, z = spheric2cartesian([0.0, arrow_dist], [np.deg2rad(ref_lat), np.deg2rad(ref_lat)], [np.deg2rad(delta_ref), np.deg2rad(delta_ref)])
+
+            # arrow plotting based on plotly hack provided through
+            # https://stackoverflow.com/a/66792953/2336056
+            arrow_tip_ratio = 0.4
+            arrow_starting_ratio = 0.95
+
+            # plot arrow line
+            fig.add_trace(go.Scatter3d(x=x,
+                                       y=y,
+                                       z=z,
+                                       mode='lines',
+                                       # marker=dict(symbol="arrow", size=15, angleref="previous", color="black"),  # only works in plotly 2d plots
+                                       name='reference long.',
+                                       showlegend=True,
+                                       line=dict(color="black", width=3),
+                                       # thetaunit="radians"
+                                       ))
+            # plot arrow head
+            fig.add_trace(go.Cone(x=[x[0] + arrow_starting_ratio*(x[1] - x[0])],
+                                  y=[y[0] + arrow_starting_ratio*(y[1] - y[0])],
+                                  z=[z[0] + arrow_starting_ratio*(z[1] - z[0])],
+                                  u=[arrow_tip_ratio*(x[1] - x[0])],
+                                  v=[arrow_tip_ratio*(y[1] - y[0])],
+                                  w=[arrow_tip_ratio*(z[1] - z[0])],
+                                  name='reference long.',
+                                  showlegend=False,
+                                  showscale=False,
+                                  colorscale=[[0, 'rgb(0,0,0)'], [1, 'rgb(0,0,0)']]
+                                  ))
+
+            if plot_spirals:
+                phi = np.ones(len(r_array))*np.deg2rad(ref_lat)
+                x, y, z = spheric2cartesian(r_array, phi, alpha_ref)
+
+                fig.add_trace(go.Scatter3d(x=x,
+                                           y=y,
+                                           z=z,
+                                           mode='lines',
+                                           name=f'field line connecting to<br>ref. long. (vsw={reference_vsw} km/s)',
+                                           showlegend=True,
+                                           line=dict(color="black", dash="dot"),
+                                           # thetaunit="radians"
+                                           ))
+            #     ax.plot(alpha_ref, r_array * np.cos(np.deg2rad(ref_lat)), '--k', label=f'field line connecting to\nref. long. (vsw={reference_vsw} km/s)')
+
+        if plot_equatorial_plane:
+            fig.add_trace(go.Surface(x=np.linspace(-200, 200, 100),
+                                    y=np.linspace(-200, 200, 100),
+                                    z=np.zeros((100, 100)),
+                                    hoverinfo='skip',
+                                    colorscale='gray', showscale=False, opacity=0.2))
 
         if _isstreamlit():
             fig.update_layout(width=700, height=700)