diff --git a/Documentation/sphinx/magnetopause.rst b/Documentation/sphinx/magnetopause.rst new file mode 100644 index 00000000..b1a225ac --- /dev/null +++ b/Documentation/sphinx/magnetopause.rst @@ -0,0 +1,157 @@ +Magnetopause: how to find +========================= + +Note: The analysator methods here are prototypes and should be adapted and modified as needed. + +Magnetopause + +Plasma beta, beta* +------------------ + + +Modified plasma beta + +.. math:: \beta * = \dfrac{P_{th}P_{dyn}}{B^2/2\mu_0} + +[Xu_et_al_2016]_, [Brenner_et_al_2021]_ + + +The beta* gets values below 1 inside the magnetosphere near magnetopause and can be used to create a magnetopause surface. + +Caveats: magnetotail current sheet has beta* :math:`>` 1 + +**in analysator:** + +see :func:`magnetopause.magnetopause` "method" keyword options "beta_star", "beta_star_with_connectivity" + +datareducer: beta_star, vg_beta_star + +.. [Xu_et_al_2016] Xu, S., M. W. Liemohn, C. Dong, D. L. Mitchell, S. W. Bougher, and Y. Ma (2016), Pressure and ion composition boundaries at Mars, J. Geophys. Res. Space Physics, 121, 6417–6429, doi:10.1002/2016JA022644. +.. [Brenner_et_al_2021] Brenner A, Pulkkinen TI, Al Shidi Q and Toth G (2021) Stormtime Energetics: Energy Transport Across the Magnetopause in a Global MHD Simulation. Front. Astron. Space Sci. 8:756732. doi: 10.3389/fspas.2021.756732 + + + +Field line connectivity +----------------------- + +Bulkfiles from newer Vlasiator runs include connection (whether magnetic field lines are closed-closed, closed-open, open/closed or open-open) and coordinates for field line start and end points that can be used in finding the mangetopause location. + + +Solar wind flow +--------------- + +Method used in e.g. [Palmroth_et_al_2003]_ + +Streamlines of velocity field that are traced from outside the bow shock curve around the magnetopause. + +Caveats: sometimes some streamlines can curve into the magnetotail or dayside magnetoshpere + + +**In analysator:** + +see +:func:`calculations.find_magnetopause_sw_streamline_2d` +:func:`calculations.find_magnetopause_sw_streamline_3d` + +Streamlines are traced from outside the bow shock towards Earth. A subsolar point for the magnetopause is chosen to be where streamlines get closest to Earth in x-axis [y/z~0]. + +From subsolar point towards Earth, space is divided radially by spherical coordiante (theta from x-axis) angles theta and phi, and magnetopause is located by looking at streamline point distances from origo and marked to be at middle of the sector + +From the Earth towards negative x the space is divided into yz-planes and streamlines are interpolated to certain x-points. +Each yz-plane is then divided into 2d sectors and magnetopause is marked to be in the middle of the sector with the radius of the n:th closest streamline to the x-axis. + + +For subsolar point, radial dayside, and -x planes the closest streamline point can be changed to be n:th closest by setting keyword *ignore*, in which case *ignore* closest streamline points are not taken into account. + + +2d: +Note: no radial dayside, uses yt-package instead of analysator's fieldtracer + +3d: +After the magnetopause points are chosen, they are made into a surface by setting connnection lines between vertices (magnetopause points) to make surface triangles. +This surface is then made into a vtkPolyData and returned as vtkDataSetSurfaceFilter that can be written into e.g. .vtp file. + + + +.. [Palmroth_et_al_2003] Palmroth, M., T. I. Pulkkinen, P. Janhunen, and C.-C. Wu (2003), Stormtime energy transfer in global MHD simulation, J. Geophys. Res., 108, 1048, doi:10.1029/2002JA009446, A1. + + + + +Shue et al. (1997) +------------------ + +The Shue et al. (1997) [Shue_et_al_1997]_ magnetopause model: + +.. math:: + + r_0 &= (11.4 + 0.013 B_z)(D_p)^{-\frac{1}{6.6}}, for B_z \geq 0 \\ + &= (11.4 + 0.14 B_z)(D_p)^{-\frac{1}{6.6}}, for B_z \leq 0 + +.. math:: + + \alpha = (0.58-0.010B_z)(1+0.010 D_p) + + +where :math:`D_p` is the dynamic pressure of the solar wind and :math:`B_z` the magnetic field z-component magnitude. +:math:`r_0` is the magnetopause standoff distance and :math:`\alpha` is the level of tail flaring. + + +The magnetopause radius as a function of angle :math:`\theta` is + +.. math:: + r = r_0 (\frac{2}{1+\cos \theta})^\alpha + +**In analysator:** + +see :doc:`../shue` + + +.. [Shue_et_al_1997] Shue, J.-H., Chao, J. K., Fu, H. C., Russell, C. T., Song, P., Khurana, K. K., and Singer, H. J. (1997). A new functional form to study the solar wind control of the magnetopause size and shape. Journal of Geophysical Research: Space Physics, 102(A5):9497–9511. eprint: https://agupubs.onlinelibrary.wiley.com/doi/pdf/10.1029/97JA00196. + + + + + +**In analysator:** +------------------ + +:mod:`magnetopause` in scripts for 3d runs +Constructs the magnetopause surface with vtk's vtkDelaunay3d triangulation with optional alpha to make the surface non-convex. +Uses regions.py functions. + +Important: SDF of non-convex surface might not always work as expected inside the magnetosphere if e.g. the lobes are hollow. + +options (magnetopause() method keyword) and some notes: + +* solar wind flow ("streamlines") + * uses *magnetopause_sw_streamline_3d.py* from pyCalculations + * if streamlines make a turn so that the velocity points sunwards (to pos. x), the streamline is ignored from that point onwards + * streamlines that enter an area where beta* is below 0.4 are also stopped + * this can affect the far magnetotail as streamlines may turn inwards into lobes at some point and may result in magnetopause not forming + due to not enough streamline points where magnetopause is supposed to be + * very dependent on how solar wind streamlines behave + * streamline seeds and other options can greatly affect the resulting magnetopause + +* beta* ("beta_star") + * beta* treshold might need tweaking as sometimes there are small low beta* areas in the magnetosheath that get taken in distorting the magnetopause shape at nose + * convex hull (Delaunay_alpha=None) usually makes a nice rough magnetopause but goes over any inward dips (like polar cusps) + * alpha shape (Delaunay_alpha= e.g. 1*R_E) does a better job at cusps and delicate shapes like vortices but might fail at SDF inside the magnetosphere + * Delaynay3d has an easier time if the treshold is something like [0.4, 0.5] and not [0.0, 0.5], but this can affect the alpha shape if alpha is used + +* beta* with magnetic field line connectivity ("beta_star_with_connectivity") + * includes closed-closed magnetic field line areas if available, otherwise like "beta_star" + * can help with nose shape as beta* can be set lower to exclude magnetosheath low beta* areas while still getting the full dayside from field lines + +* Shue et al. 1997 ("shue") + * uses *shue.py* from scripts + * a rough theoretical magnetopause using Shue et al. 1997 method based on B_z, solar wind density, and solar wind velocity + +* user-defined parameter tresholds ("dict") + * creates a magnetopause (or other area) using the Delaunay3d triangulation of some area where user-defined tresholds given as dictionary + * dictionary key is data name in datafile and value is treshold used, if dictionary has multiple conditions, they all need to be fulfilled + * dictionary example: {"vg_rho": [None, 1e5]} makes a magnetopause using cells where density is less than 1e5 + + +.. automodule:: magnetopause + :members: \ No newline at end of file diff --git a/Documentation/sphinx/magnetopause2d.rst b/Documentation/sphinx/magnetopause2d.rst deleted file mode 100644 index 8e1717f1..00000000 --- a/Documentation/sphinx/magnetopause2d.rst +++ /dev/null @@ -1,5 +0,0 @@ -magnetopause2d --------------- - -.. automodule:: magnetopause2d - :members: \ No newline at end of file diff --git a/Documentation/sphinx/magnetopause3d.rst b/Documentation/sphinx/magnetopause3d.rst deleted file mode 100644 index 7d90752f..00000000 --- a/Documentation/sphinx/magnetopause3d.rst +++ /dev/null @@ -1,5 +0,0 @@ -magnetopause3d --------------- - -.. automodule:: magnetopause3d - :members: \ No newline at end of file diff --git a/Documentation/sphinx/magnetosphere_regions.rst b/Documentation/sphinx/magnetosphere_regions.rst new file mode 100644 index 00000000..a4e7fa2d --- /dev/null +++ b/Documentation/sphinx/magnetosphere_regions.rst @@ -0,0 +1,228 @@ +Magnetosphere regions and bow shock: How to find +================================================ + +Note: The analysator methods here are prototypes and should be adapted and modified as needed. + +Bow shock +--------- + +Plasma properties for estimating bow shock position: + +* plasma compression: + * :math:`n_p > 2n_{p, sw}` [Battarbee_et_al_2020]_ (Vlasiator) +* solar wind core heating: + * :math:`T_{core} > 4T_{sw}` [Battarbee_et_al_2020]_ (Vlasiator) + * :math:`T_{core} = 3T_{sw}` [Suni_et_al_2021]_ (Vlasiator) +* magnetosonic Mach number: + * :math:`M_{ms} < 1` [Battarbee_et_al_2020]_ (Vlasiator) + + +**In analysator:** + +:mod:`regions` has an option to find the bow shock. Default method uses 1.5*solar wind density as limit. +Usage example: + +.. code-block:: python + + datafile = "vlsvbulkfile.vlsv" + outfilen = "bowshock.vlsv" + RegionFlags(datafile, outfilen, regions=["bowshock"]) + + +Magnetosheath +------------- + +properties: + +* density: + * :math:`8 cm^{-3}` [Hudges_Introduction_to_space_physics_Ch_9]_ +* temperature: + * ion: :math:`150 eV` [Hudges_Introduction_to_space_physics_Ch_9]_ + * electron: :math:`25 eV` [Hudges_Introduction_to_space_physics_Ch_9]_ +* magnetic field: + * :math:`15 nT` [Hudges_Introduction_to_space_physics_Ch_9]_ +* plasma :math:`\beta`: + * 2.5 [Hudges_Introduction_to_space_physics_Ch_9]_ + +**In analysator:** + +:mod:`regions` has an option to find the magnetosheath using bow shock and magnetopause: +Usage example: + +.. code-block:: python + + datafile = "vlsvbulkfile.vlsv" + outfilen = "magnetosheath.vlsv" + RegionFlags(datafile, outfilen, regions=["magnetosheath"]) + + +Polar cusps +----------- + +*Properties:* + +* plasma density: high density in comparison to solar wind + * Ion density :math:`\geq` solar wind ion density [Pitout_et_al_2006]_ (Cluster spacecraft data) +* ion energy: + * mean ion energy :math:`~2-3 keV` [Pitout_et_al_2006]_ /[Stenuit_et_al_2001]_ +* energy flux: + * energy flux + + + +**In analysator:** + +:mod:`regions` has an option to find cusps using convex hull of the magnetosphere. +Note that the default parameters may catch some dayside cells near the magnetopause and some non-cusp areas inside dayside in addition to cusps. +Areas near magnetopause can be excluded for example by only including cells where the SDF is less than e.g. -1 Re. + +Usage example: + +.. code-block:: python + + datafile = "vlsvbulkfile.vlsv" + outfilen = "cusps.vlsv" + RegionFlags(datafile, outfilen, regions=["cusps"]) + + +Tail lobes +---------- + +* plasma density: low + * below :math:`0.03 cm^{-3}` [Grison_et_al_2025]_ (Cluster spacecraft data) + * :math:`0.01 cm^{-3}` [Koskinen_Space_Storms]_ p.38 + * less than :math:`0.1 cm^{-3}` [Wolf_Introduction_to_space_physics_Ch_10]_ p.291 +* plasma :math:`\beta`: low + * typically around :math:`0.05` [Grison_et_al_2025]_ (Cluster spacecraft data) + * :math:`3e-3` [Koskinen_Space_Storms]_ p.38 +* temperature: + * ion temperature :math:`300 eV` [Koskinen_Space_Storms]_ p.38 + * electron temperature :math:`50 eV` [Koskinen_Space_Storms]_ p.38 +* magnetic field: + * :math:`20 nT` [Koskinen_Space_Storms]_ p.38 +* open magnetic field lines [Wolf_Introduction_to_space_physics_Ch_10]_ p.291 +* strong and stable magnetic field towards the Earth (northern lobe) and away from the Earth (southern lobe) [Coxon_et_al_2016]_ + +Separated from the plasma sheet by the plasma sheet boundary layer (PSBL) + + +**In analysator:** + +:mod:`regions` has an option to find tail lobes. +Note that the default parameters to find lobes may include some areas from the dayside. + +Usage example: + +.. code-block:: python + + datafile = "vlsvbulkfile.vlsv" + outfilen = "lobes.vlsv" + RegionFlags(datafile, outfilen, regions=["lobes"]) + + + + +Low-latitude boundary layer (LLBL) +---------------------------------- + + + +Properties: + +* density: + * ion number densities between those of magnetosphere and magnetosheath [Hudges_Introduction_to_space_physics_Ch_9]_ p.267 +* temperature + * ion temperatures between those of magnetosphere and magnetosheath [Hudges_Introduction_to_space_physics_Ch_9]_ p.267 +* unknown field line configuration, probably a mix of open and closed field lines [Hudges_Introduction_to_space_physics_Ch_9]_ p.262 + + + +High-latitude boundary layer (HLBL) +----------------------------------- + +Includes the plasma mantle on the tail side and the entry layer on the dayside + +Properties: + +* open magnetic field lines [Hudges_Introduction_to_space_physics_Ch_9]_ p.261 + + + + + +Plasma sheet boundary layer (PSBL) +---------------------------------- + +The plasma sheet boundary layer is a very thin boundary layer separating the tail lobes from the tail plasma sheet [Koskinen_Johdatus]_ + +*Properties:* + +* density: + * :math:`0.1 cm^{-3}` [Koskinen_Space_Storms]_ p.38 +* temperature: + * ion temperature :math:`1000 eV` [Koskinen_Space_Storms]_ p.38 + * electron temperature :math:`150 eV` [Koskinen_Space_Storms]_ p.38 +* magnetic field: + * :math:`20 nT` [Koskinen_Space_Storms]_ p.38 +* plasma :math:`\beta` : + * :math:`0.1` [Koskinen_Space_Storms]_ p.38 +* probably closed magnetic field lines [Wolf_Introduction_to_space_physics_Ch_10]_ p.291 + + + + +Central plasma sheet +-------------------- + + +*Properties:* + +* density: + * :math:`0.3 cm^{-3}` [Koskinen_Space_Storms]_ p.38 + * :math:`0.1-1 cm^{-3}` [Wolf_Introduction_to_space_physics_Ch_10]_ p.291 +* temperature: hot + * ion temperature :math:`4200 eV` [Koskinen_Space_Storms]_ p.38 + * electron temperature :math:`600 eV` [Koskinen_Space_Storms]_ p.38 +* magnetic field: + * :math:`10 nT` [Koskinen_Space_Storms]_ p.38, [Hudges_Introduction_to_space_physics_Ch_9]_ +* plasma :math:`\beta`: high + * :math:`6` [Koskinen_Space_Storms]_ p.38 +* Mostly closed magnetic field lines [Wolf_Introduction_to_space_physics_Ch_10]_ + +Inner plasma sheet: unusually low plasma beta may exist (e.g., cold tenuous plasma near the neutral sheet after long periods of northward IMF) [Boakes_et_al_2014]_, (Cluster spacecraft data) + + +**In analysator:** + +:mod:`regions` has an option to find the central plasma sheet. +Note that the default parameters may catch some dayside cells in addition to plasma sheet. +Usage example: + +.. code-block:: python + + datafile = "vlsvbulkfile.vlsv" + outfilen = "CPS.vlsv" + RegionFlags(datafile, outfilen, regions=["central_plasma_sheet"]) + + + + +.. automodule:: regions + :members: + + +------------ + +References + +.. [Battarbee_et_al_2020] Battarbee, M., Ganse, U., Pfau-Kempf, Y., Turc, L., Brito, T., Grandin, M., Koskela, T., and Palmroth, M.: Non-locality of Earth's quasi-parallel bow shock: injection of thermal protons in a hybrid-Vlasov simulation, Ann. Geophys., 38, 625-643, https://doi.org/10.5194/angeo-38-625-2020, 2020 +.. [Suni_et_al_2021] Suni, J., Palmroth, M., Turc, L., Battarbee, M., Johlander, A., Tarvus, V., et al. (2021). Connection between foreshock structures and the generation of magnetosheath jets: Vlasiator results. Geophysical Research Letters, 48, e2021GL095655. https://doi. org/10.1029/2021GL095655 +.. [Grison_et_al_2025] Grison, B., Darrouzet, F., Maggiolo, R. et al. Localization of the Cluster satellites in the geospace environment. Sci Data 12, 327 (2025). https://doi.org/10.1038/s41597-025-04639-z +.. [Koskinen_Johdatus] Koskinen, H. E. J. (2011). Johdatus plasmafysiikkaan ja sen avaruussovellutuksiin. Limes ry. +.. [Koskinen_Space_Storms] Koskinen, H. E. J. (2011). Physics of Space Storms: From the Solar Surface to the Earth. Springer-Verlag. https://doi.org/10.1007/978-3-642-00319-6 +.. [Pitout_et_al_2006] Pitout, F., Escoubet, C. P., Klecker, B., and Rème, H.: Cluster survey of the mid-altitude cusp: 1. size, location, and dynamics, Ann. Geophys., 24, 3011–3026, https://doi.org/10.5194/angeo-24-3011-2006, 2006. +.. [Coxon_et_al_2016] Coxon,J.C.,C.M.Jackman, M. P. Freeman, C. Forsyth, and I. J. Rae (2016), Identifying the magnetotail lobes with Cluster magnetometer data, J. Geophys. Res. Space Physics, 121, 1436–1446, doi:10.1002/2015JA022020. +.. [Hudges_Introduction_to_space_physics_Ch_9] Hudges, W. J. (1995) The magnetopause, magnetotail and magnetic reconnection. In Kivelson, M. G., & Russell, C. T. (Eds.), Introduction to space physics (pp.227-287). Cambridge University Press. +.. [Wolf_Introduction_to_space_physics_Ch_10] Wolf, R. A. (1995) Magnetospheric configuration. In Kivelson, M. G., & Russell, C. T. (Eds.), Introduction to space physics (pp.288-329). Cambridge University Press. +.. [Boakes_et_al_2014] Boakes, P. D., Nakamura, R., Volwerk, M., and Milan, S. E. (2014). ECLAT Cluster Spacecraft Magnetotail Plasma Region Identifications (2001–2009). Dataset Papers in Science, 2014(1):684305. eprint: https://onlinelibrary.wiley.com/doi/pdf/10.1155/2014/684305 +.. [Stenuit_et_al_2001] Stenuit, H., Sauvaud, J.-A., Delcourt, D. C., Mukai, T., Kokubun, S., Fujimoto, M., Buzulukova, N. Y., Kovrazhkin, R. A., Lin, R. P., and Lepping, R. P. (2001). A study of ion injections at the dawn and dusk polar edges of the auroral oval. Journal of Geophysical Research: Space Physics, 106(A12):29619–29631. eprint: https://agupubs.onlinelibrary.wiley.com/doi/pdf/10.1029/2001JA900060. diff --git a/Documentation/sphinx/plot_streamline_magnetopause_2d.rst b/Documentation/sphinx/plot_streamline_magnetopause_2d.rst new file mode 100644 index 00000000..38e7d4fc --- /dev/null +++ b/Documentation/sphinx/plot_streamline_magnetopause_2d.rst @@ -0,0 +1,6 @@ +plot_streamline_magnetopause_2d +------------------------------- + +.. automodule:: plot_streamline_magnetopause_2d + :members: + \ No newline at end of file diff --git a/Documentation/sphinx/scripts.rst b/Documentation/sphinx/scripts.rst index ebb57f90..f9a7d8bd 100644 --- a/Documentation/sphinx/scripts.rst +++ b/Documentation/sphinx/scripts.rst @@ -30,21 +30,12 @@ gics ------------ -magnetopause2d --------------- -:doc:`magnetopause2d` - -.. automodule:: magnetopause2d - :no-index: - ------------- +plot_streamline_magnetopause_2d +------------------------------- +:doc:`plot_streamline_magnetopause_2d` -magnetopause3d --------------- -:doc:`magnetopause3d` - -.. automodule:: magnetopause3d - :no-index: +.. automodule:: plot_streamline_magnetopause_2d + :no-index: ------------ @@ -84,6 +75,24 @@ tsyganenko ------------ +magnetopause +------------ +:doc:`magnetopause` + +.. automodule:: magnetopause + :no-index: + +------------ + +regions +------- +:doc:`magnetosphere_regions` + +.. automodule:: regions + :no-index: + +------------ + .. toctree:: :maxdepth: 2 :caption: Scripts: @@ -91,10 +100,11 @@ tsyganenko biot_savart cutthrough_timeseries gics - magnetopause2d - magnetopause3d + plot_streamline_magnetopause_2d obliqueshock obliqueshock_nif shue tsyganenko + magnetopause + magnetosphere_regions diff --git a/analysator/pyCalculations/calculations.py b/analysator/pyCalculations/calculations.py index 7a768272..3ce6790f 100644 --- a/analysator/pyCalculations/calculations.py +++ b/analysator/pyCalculations/calculations.py @@ -61,3 +61,5 @@ from non_maxwellianity import epsilon_M from null_lines import LMN_null_lines_FOTE from interpolator_amr import AMRInterpolator, supported_amr_interpolators +from magnetopause_sw_streamline_2d import find_magnetopause_sw_streamline_2d +from magnetopause_sw_streamline_3d import find_magnetopause_sw_streamline_3d diff --git a/analysator/pyCalculations/fieldtracer.py b/analysator/pyCalculations/fieldtracer.py index aa8d521c..e734a25d 100644 --- a/analysator/pyCalculations/fieldtracer.py +++ b/analysator/pyCalculations/fieldtracer.py @@ -51,7 +51,7 @@ def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar ''' Field tracer in a static frame :param vlsvReader: An open vlsv file - :param x: Starting point for the field trace + :param x0: Starting point for the field trace :param max_iterations: The maximum amount of iteractions before the algorithm stops :param dx: One iteration step length :param direction: '+' or '-' or '+-' Follow field in the plus direction or minus direction @@ -61,6 +61,10 @@ def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar :returns: List of coordinates ''' + # Check the starting point + if (np.shape(x0) != (3,)): + raise ValueError("Starting point should be a single [x,y,z] coordinate in 1-dimensional list or array") + if(bvar != 'B'): warnings.warn("User defined tracing variable detected. fg, volumetric variable results may not work as intended, use face-values instead.") @@ -117,7 +121,7 @@ def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar x = np.arange(mins[0], maxs[0], dcell[0]) + 0.5*dcell[0] y = np.arange(mins[1], maxs[1], dcell[1]) + 0.5*dcell[1] z = np.arange(mins[2], maxs[2], dcell[2]) + 0.5*dcell[2] - coordinates = np.array([x,y,z]) + coordinates = [x,y,z] # Debug: assert len(x) == sizes[0] @@ -141,7 +145,7 @@ def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar x = np.arange(mins[0], maxs[0], dcell[0]) + 0.5*dcell[0] y = np.arange(mins[1], maxs[1], dcell[1]) + 0.5*dcell[1] z = np.arange(mins[2], maxs[2], dcell[2]) + 0.5*dcell[2] - coordinates = np.array([x,y,z]) + coordinates = [x,y,z] # Debug: assert len(x) == sizes[0] @@ -284,7 +288,7 @@ def find_unit_vector(vg, coord): points_traced_unique[mask_update,i,:] = next_points[mask_update,:] # distances = np.linalg.norm(points_traced_unique[:,i,:],axis = 1) - mask_update[stop_condition(vlsvReader, points_traced_unique[:,i,:])] = False + mask_update[stop_condition(vlsvReader, points_traced_unique[:,i,:], vlsvReader.read_interpolated_variable(vg,points_traced_unique[:,i,:]))] = False # points_traced_unique[~mask_update, i, :] = points_traced_unique[~mask_update, i-1, :] @@ -292,7 +296,7 @@ def find_unit_vector(vg, coord): return points_traced # Default stop tracing condition for the vg tracing, (No stop until max_iteration) -def default_stopping_condition(vlsvReader, points): +def default_stopping_condition(vlsvReader, points, vars): [xmin, ymin, zmin, xmax, ymax, zmax] = vlsvReader.get_spatial_mesh_extent() x = points[:, 0] y = points[:, 1] @@ -353,6 +357,10 @@ def static_field_tracer_3d( vlsvReader, seed_coords, max_iterations, dx, directi # Standardize input: (N,3) np.array if type(seed_coords) != np.ndarray: raise TypeError("Please give a numpy array.") + + # Transform a single (3,) coordinate to (1,3) if needed + if np.shape(seed_coords)==(3,): + seed_coords = np.array([seed_coords]) # Cache and read variables: vg = None diff --git a/analysator/pyCalculations/magnetopause_sw_streamline_2d.py b/analysator/pyCalculations/magnetopause_sw_streamline_2d.py new file mode 100644 index 00000000..4cf2a3b2 --- /dev/null +++ b/analysator/pyCalculations/magnetopause_sw_streamline_2d.py @@ -0,0 +1,162 @@ +''' +Finds the magnetopause position by tracing streamlines of the plasma flow for two-dimensional Vlasiator runs. Needs the yt package. +''' + +import numpy as np +import analysator as pt +import yt + + +def interpolate(streamline, x_points): + """Interpolates a single streamline for make_magnetopause(). + + :param streamline: a single streamline to be interpolated + :param x_points: points in the x-axis to use for interpolation + :returns: the streamline as numpy array of x,z coordinate points where the x-axis coordinates are the points given to the function + """ + + arr = np.array(streamline) + + # set arrays for interpolation + xp = arr[:,0][::-1] + zp = arr[:,2][::-1] + + #interpolate z coordinates + z_points = np.interp(x_points, xp, zp, left=np.NaN, right=np.NaN) + + return np.array([x_points, z_points]) + + +def make_streamlines(vlsvfile, streamline_seeds=None,seeds_n=200, seeds_x0=20*6371000, seeds_range=[-5*6371000, 5*6371000], streamline_length=40*6371000): + """Traces streamlines of velocity field using the yt package. + + :param vlsvfile: directory and file name of .vlsv data file to use for VlsvReader + :kword streamline_seeds: optional streamline starting points in numpy array (coordinates in meters including the y-coordinate 0.0) + :kword seeds_n: instead of streamline_seeds provide a number of streamlines to be traced + :kword seeds_x0: instead of streamline_seeds provide an x-coordinate for streamline starting points + :kword seeds_range: instead of streamline_seeds provide [min, max] range to use for streamline starting point z-coordinates + :kword streamline_length: streamline length + + :returns: streamlines as numpy array + """ + + # bulk file + f = pt.vlsvfile.VlsvReader(file_name=vlsvfile) + + # get box coordinates from data + [xmin, ymin, zmin, xmax, ymax, zmax] = f.get_spatial_mesh_extent() + [xsize, ysize, zsize] = f.get_spatial_mesh_size() + simext =[xmin,xmax,ymin,ymax,zmin,zmax] + sizes = np.array([xsize,ysize,zsize]) + boxcoords=list(simext) + + cellids = f.read_variable("CellID") + + #Read the data from vlsv-file + Vx = f.read_variable("v", operator="x") + Vz = f.read_variable("v", operator="z") + + + #Re-shape variable data + Vxs=Vx[np.argsort(cellids)].reshape(f.get_spatial_mesh_size(), order="F") + Vys = np.zeros_like(Vxs) + Vzs=Vz[np.argsort(cellids)].reshape(f.get_spatial_mesh_size(), order="F") + + data=dict(Vx=Vxs,Vy=Vys,Vz=Vzs) + + #Create starting points for streamlines if they are not given + if streamline_seeds == None: + streamline_seeds = np.array([[seeds_x0, 0 ,i] for i in np.linspace(seeds_range[0], seeds_range[1], seeds_n)]) + + #streamline_seeds = np.array(streamline_seeds) + #dataset in yt-form + yt_dataset = yt.load_uniform_grid( + data, + sizes, + bbox=np.array([[boxcoords[0], boxcoords[1]], + [boxcoords[2],boxcoords[3]], + [boxcoords[4],boxcoords[5]]])) + + + #data, seeds, dictionary positions, step size + streamlines = yt.visualization.api.Streamlines(yt_dataset, streamline_seeds, + "Vx", "Vy", "Vz", length=streamline_length, direction=1) + + #trace the streamlines with yt + streamlines.integrate_through_volume() + # return streamline positions + return np.array(streamlines.streamlines) + + +def make_magnetopause(streamlines, end_x=-15*6371000, x_point_n=50): + """Finds the mangetopause location based on streamlines. + + :param streams: streamlines (coordinates in m) + :kword end_x: tail end x-coordinate (how far along the negative x-axis the magnetopause is calculated) + :kword x_point_n: integer, how many x-axis points the magnetopause will be divided in between the subsolar point and tail + + :returns: the magnetopause position as coordinate points in numpy array + """ + + RE = 6371000 + + streampoints = np.reshape(streamlines, (streamlines.shape[0]*streamlines.shape[1], 3)) #all the points in one array + + ## find the subsolar dayside point in the positive x-axis + ## do this by finding a stremline point on positive x axis closest to the Earth + x_axis_points = streampoints[(streampoints[:,2]< RE) & (streampoints[:,2]> -RE) & (streampoints[:,0]> 0)] + subsolar_x =np.min(x_axis_points[:,0]) + + ## define points in the x axis where to find magnetopause points on the yz-plane + x_points = np.linspace(subsolar_x, end_x, x_point_n) + + ## interpolate more exact points for streamlines at exery x_point + new_streampoints = np.zeros((len(x_points), len(streamlines), 1)) # new array for keeping interpolated streamlines in form streamlines_new[x_point, streamline, z-coordinate] + + for i,stream in enumerate(streamlines): + interpolated_streamline = interpolate(stream, x_points) + for j in range(0, len(x_points)): + new_streampoints[j, i,:] = interpolated_streamline[1,j] + + + ## start making the magnetopause + ## in each x_point, find the closest streamline to x-axis in the positive and negative z-axis + + pos_z_mpause = np.zeros((len(x_points), 2)) + neg_z_mpause = np.zeros((len(x_points), 2)) + + for i, x_point in enumerate(x_points): + pos = new_streampoints[i, new_streampoints[i,:] > 0] + neg = new_streampoints[i, new_streampoints[i,:] < 0] + + if (pos.size == 0) or (neg.size == 0): + raise ValueError('No streamlines found for x axis point, try adding streamlines or checking the x_points') + + # find points closest to x-axis and save found points + pos_z_mpause[i] = [x_point, pos[pos.argmin()]] + neg_z_mpause[i] = [x_point, neg[neg.argmax()]] + + magnetopause = np.concatenate((pos_z_mpause[::-1], np.array([[subsolar_x, 0]]), neg_z_mpause)) + + return magnetopause + + +def find_magnetopause_sw_streamline_2d(vlsvfile, streamline_seeds=None, seeds_n=200, seeds_x0=20*6371000, seeds_range=[-5*6371000, 5*6371000], streamline_length=45*6371000, end_x=-15*6371000, x_point_n=50): + """Finds the magnetopause position by tracing streamlines of the velocity field for 2d runs. + + :param vlsvfile: directory and file name of .vlsv data file to use for VlsvReader + :kword streamline_seeds: optional streamline starting points in numpy array (coordinates in meters including the y-coordinate 0.0) + :kword seeds_n: instead of streamline_seeds provide a number of streamlines to be traced + :kword seeds_x0: instead of streamline_seeds provide an x-coordinate for streamline starting points + :kword seeds_range: instead of streamline_seeds provide [min, max] range to use for streamline starting point z-coordinates + :kword streamline_length: streamline length for tracing + :kword end_x: tail end x-coordinate (how far along the negative x-axis the magnetopause is calculated) + :kword x_point_n: integer, how many x-axis points the magnetopause will be divided in between the subsolar point and tail + + :returns: the magnetopause position as coordinate points in numpy array + """ + + streamlines = make_streamlines(vlsvfile, streamline_seeds, seeds_n, seeds_x0, seeds_range, streamline_length) + magnetopause = make_magnetopause(streamlines, end_x, x_point_n) + + return magnetopause diff --git a/analysator/pyCalculations/magnetopause_sw_streamline_3d.py b/analysator/pyCalculations/magnetopause_sw_streamline_3d.py new file mode 100644 index 00000000..46520daf --- /dev/null +++ b/analysator/pyCalculations/magnetopause_sw_streamline_3d.py @@ -0,0 +1,468 @@ +''' +Finds the magnetopause position by tracing streamlines of the plasma flow for three-dimensional Vlasiator runs. +''' + +import matplotlib.pyplot as plt +import numpy as np +import analysator as pt +import vtk + +def cartesian_to_polar(cartesian_coords): # for segments of plane + """Converts cartesian coordinates to polar (for the segments of the yz-planes). + + :param cartesian_coords: (y,z) coordinates as list or array + :returns: the polar coordinates r, phi (angle in degrees) + """ + y,z = cartesian_coords[0], cartesian_coords[1] + r = np.sqrt(z**2 + y**2) + phi = np.arctan2(z, y) + phi = np.rad2deg(phi) #angle in degrees + if phi < 0: phi = phi+360 + return(r, phi) + +def polar_to_cartesian(r, phi): + """Converts polar coordinates of the yz-plane to cartesian coordinates. + + :param r: radius of the segment (distance from the x-axis) + :param phi: the angle coordinate in degrees + :returns: y, z -coordinates in cartesian system + """ + phi = np.deg2rad(phi) + y = r * np.cos(phi) + z = r * np.sin(phi) + return(y, z) + + +def cartesian_to_spherical(cartesian_coords): + """ Cartesian to spherical coordinates (Note: axes are swapped!)""" + x, y, z = cartesian_coords[0], cartesian_coords[1], cartesian_coords[2] + r = np.sqrt(x**2+y**2+z**2) + theta = np.arctan2(np.sqrt(y**2+z**2), x) #inclination + phi = np.arctan2(z, y)+np.pi # azimuth + + return [r, theta, phi] + + + +def spherical_to_cartesian(sphe_coords): + r, theta, phi = sphe_coords[0], sphe_coords[1], sphe_coords[2] + y = r*np.sin(theta)*np.cos(phi) + z = r*np.sin(theta)*np.sin(phi) + x = r*np.cos(theta) + return [x, y, z] + + +def interpolate(streamline, x_points): + """Interpolates a single streamline for make_magnetopause(). + + :param streamline: a single streamline to be interpolated + :param x_points: points in the x-axis to use for interpolation + :returns: the streamline as numpy array of coordinate points where the x-axis coordinates are the points given to the function + """ + arr = np.array(streamline) + + # set arrays for interpolation + xp = arr[:,0][::-1] + yp = arr[:,1][::-1] + zp = arr[:,2][::-1] + + #interpolate z from xz-plane + z_points = np.interp(x_points, xp, zp, left=np.nan, right=np.nan) + #interpolate y from xy-plane + y_points = np.interp(x_points, xp, yp, left=np.nan, right=np.nan) + + + return np.array([x_points, y_points, z_points]) + + + + +def make_surface(coords): + '''Defines a surface constructed of input coordinates as triangles. + + :param coords: points that make the surface + :returns: list of verts and vert indices of surface triangles as numpy arrays + + input coordinates must be in form [...[c11, c21, c31, ... cn1],[c12, c22, c32, ... cn2],... + where cij = [xij, yij, zij], i marks sector, j marks yz-plane (x_coord) index + + Three points make a triangle, triangles make the surface. + For every two planes next to each other: + - take every other point from plane1, every other from plane2 (in order!) + - from list of points: every three points closest to each other make a surface + + Example: + plane 1: [v1, v2, v3, v4] + plane 2: [v5, v6, v7, v8] + + -> list: [v1, v5, v2, v6, v3,...] + -> triangles: + v1 v5 v2 + v5 v2 v6 + v2 v6 v3 + . + . + . + + ''' + verts = [] #points + faces = [] #surface triangles + + slices_in_plane = len(coords[0]) + planes = len(coords) + + #get points + for plane in coords: + for vert in plane: + verts.append(vert) + + #get triangle (face) indices + #Let's consider the area between the first two planes + #ind1, ind2, ind3 for triangle indices + ind1 = 0 #starts from plane 1 + ind2 = slices_in_plane #starts from plane 2 + ind3 = 1 #starts from plane 1 + first_triangles = [] + + while len(first_triangles) < (2*slices_in_plane): + first_triangles.append([ind1, ind2, ind3]) + ind1 = ind2 + ind2 = ind3 + if (ind3 == (slices_in_plane*2)-1): #max index, go over to plane 1 first index + ind3 = 0 + elif (ind3 == 0): #last round, go to plane 2 first index + ind3 = slices_in_plane + else: + ind3 = ind1 + 1 + + first_triangles = np.array(first_triangles) + + #Now the rest of the triangles are just the first triangles + (index of area * slices_in_plane) + for area_index in range(planes-1): + next_triangles = [x + slices_in_plane*area_index for x in first_triangles] + faces.extend(next_triangles) + + #### test area #### + # From last triangles remove every other triangle + # (a single subsolar point -> last triangles are actual triangles instead of rectangles sliced in two) + #removed = 0 + #for i in range(len(faces)-slices_in_plane*2, len(faces)): + # if i%2!=0: + # faces.pop(i-removed) + # removed += 1 + + # From last triangles remove every other triangle + # (a single subsolar point -> last triangles are actual triangles instead of rectangles sliced in two) + # Also fix the last triangles so that they only point to one subsolar point and have normals towards outside + #subsolar_index = int(len(verts)-slices_in_plane) + + #for i,triangle in enumerate(reversed(faces)): + # if i > (slices_in_plane): # faces not in last plane (we're going backwards) + # break + + # faces[len(faces)-i-1] = np.clip(triangle, a_min=0, a_max=subsolar_index) + + # this would remove duplicate subsolar points from vertices but makes 2d slicing harder + #verts = verts[:int(len(verts)-slices_in_plane+1)] + + # Change every other face triangle (except for last slice triangles) normal direction so that all face the same way (hopefully) + #faces = np.array(faces) + #for i in range(len(faces)-int(slices_in_plane)): + # if i%2!=0: + # faces[i,1], faces[i,2] = faces[i,2], faces[i,1] + + ################### + + # Change every other face triangle normal direction so that all face the same way + faces = np.array(faces) + for i in range(len(faces)): + if i%2==0: + faces[i,1], faces[i,2] = faces[i,2], faces[i,1] + + return np.array(verts), faces + + +def make_vtk_surface(vertices, faces): + """Makes a vtk DataSetSurfaceFilter from vertices and faces of a triangulated surface. + + :param vertices: vertex points of triangles in shape [[x0, y0, z0], [x1, y1, z1],...] + :param faces: face connectivities as indices of vertices so that triangle normals point outwards + :returns: a vtkDataSetSurfaceFilter object + """ + + points = vtk.vtkPoints() + + for vert in vertices: + points.InsertNextPoint(vert) + + # make vtk PolyData object + + triangles = vtk.vtkCellArray() + for face in faces: + triangle = vtk.vtkTriangle() + triangle.GetPointIds().SetId(0, face[0]) + triangle.GetPointIds().SetId(1, face[1]) + triangle.GetPointIds().SetId(2, face[2]) + triangles.InsertNextCell(triangle) + + polydata = vtk.vtkPolyData() + polydata.SetPoints(points) + polydata.Modified() + polydata.SetPolys(triangles) + polydata.Modified() + + surface = vtk.vtkDataSetSurfaceFilter() + surface.SetInputData(polydata) + surface.Update() + + return surface + +def streamline_stopping_condition(vlsvReader, points, value): + [xmin, ymin, zmin, xmax, ymax, zmax] = vlsvReader.get_spatial_mesh_extent() + x = points[:, 0] + y = points[:, 1] + z = points[:, 2] + beta_star = vlsvReader.read_interpolated_variable("vg_beta_star", points) + return (x < xmin)|(x > xmax) | (y < ymin)|(y > ymax) | (z < zmin)|(z > zmax)|(value[:,0] > 0)| (beta_star < 0.4) + + +def make_streamlines(vlsvfile, streamline_seeds=None, seeds_n=25, seeds_x0=20*6371000, seeds_range=[-5*6371000, 5*6371000], dl=2e6, iterations=200): + """Traces streamlines of velocity field from outside the magnetosphere to magnetotail. + Stopping condition for when streamlines turn sunwards, go out of box, or hit beta* below 0.4 region + + :param vlsvfile: directory and file name of .vlsv data file to use for VlsvReader + + :kword streamline_seeds: optional streamline starting points in numpy array + :kword seeds_n: instead of streamline_seeds provide a number of streamlines to be traced + :kword seeds_x0: instead of streamline_seeds provide an x-coordinate for streamline starting points + :kword seeds_range: instead of streamline_seeds provide [min, max] range to use for streamline starting point coordinates (both y- and z-directions use the same range) + :kword dl: streamline iteration step length in m + :kword iterations: int, number of iteration steps + + :returns: streamlines as numpy array + """ + + f = pt.vlsvfile.VlsvReader(file_name=vlsvfile) + + # Create streamline starting points if needed + if streamline_seeds == None: + streamline_seeds = np.zeros([seeds_n**2, 3]) + + t = np.linspace(seeds_range[0], seeds_range[1], seeds_n) + k = 0 + for i in t: + for j in t: + streamline_seeds[k] = [seeds_x0, i, j] + k = k+1 + + # Trace the streamlines + streams = pt.calculations.static_field_tracer_3d( + vlsvReader=f, + seed_coords=streamline_seeds, + max_iterations=iterations, + dx=dl, + direction='+', + grid_var='vg_v', + stop_condition=streamline_stopping_condition + ) + + return streams + + +def make_magnetopause(streams, end_x=-15*6371000, x_point_n=50, sector_n=36, ignore=0): + """Finds the magnetopause location based on streamlines. + + :param streams: streamlines (coordinates in m) + + :kword end_x: tail end x-coordinate (how far along the negative x-axis the magnetopause is calculated) + :kword x_point_n: integer, how many x-axis points the magnetopause will be divided in between the subsolar point and tail + :kword sector_n: integer, how many sectors the magnetopause will be divided in on each yz-plane + + :returns: the magnetopause position as coordinate points in numpy array, form [...[c11, c21, c31, ... cn1],[c12, c22, c32, ... cn2],... + where cij = [xij, yij, zij], i marks sector, j marks yz-plane (x-coordinate) index + """ + + RE = 6371000 + + ## if given sector number isn't divisible by 4, make it so because we want to have magnetopause points at exactly y=0 and z=0 for 2d slices of the whole thing + while sector_n%4 != 0: + sector_n +=1 + + #streams = streams*(1/RE) # streamlines in rE + streampoints = np.reshape(streams, (streams.shape[0]*streams.shape[1], 3)) #all the points in one array) + + ## find the subsolar dayside point in the x-axis + ## do this by finding a streamline point on positive x axis closest to the Earth + # streampoints closer than ~1 rE to positive x-axis: + x_axis_points = streampoints[(streampoints[:,1]-RE) & (streampoints[:,2]>-RE) & (streampoints[:,0]>0) & (streampoints[:,0]>0)] + if ignore == 0: + subsolar_x =np.min(x_axis_points[:,0]) + else: + subsolar_x = np.partition(x_axis_points[:,0], ignore)[ignore] # take the nth point as subsolar point + + # divide the x point numbers between x > 0 (radial) an x < 0 (yz-planes) by ratio + dayside_x_point_n = int((subsolar_x/np.abs(end_x))*x_point_n) + + ### dayside magnetopause ### + # for x > 0, look for magnetopause radially + dayside_points = streampoints[streampoints[:,0] > 0] + + phi_slices = sector_n + theta_slices = dayside_x_point_n + phi_step = 2*np.pi/phi_slices + theta_step = np.pi/(2*theta_slices) # positive x-axis only + + def cartesian_to_spherical_grid(cartesian_coords): + # Only for x > 0 + r, theta, phi = cartesian_to_spherical(cartesian_coords) + theta_idx = int(theta/theta_step) + #phi_idx = int(phi/phi_step) # off by half grid cell in relation to x<0 grid + if (phi (0-phi_step*0.5)): + phi_idx = 0 + else: + phi_idx = round(phi/phi_step) + + return [theta_idx, phi_idx, r] + + def grid_mid_point(theta_idx, phi_idx): + return (theta_idx+0.5)*theta_step, (phi_idx)*phi_step + + # make a dictionary based on spherical areas by theta index and phi index + sph_points = {} + for point in dayside_points: + sph_gridpoint = cartesian_to_spherical_grid(point) + idxs = (sph_gridpoint[0],sph_gridpoint[1]) # key is (theta_i, phi_i) + if idxs not in sph_points: + sph_points[idxs] = [sph_gridpoint[2]] + else: + sph_points[idxs].append(sph_gridpoint[2]) + + # dayside magetopause from subsolar point towards origo + dayside_magnetopause = np.zeros((theta_slices, phi_slices, 3)) + for ring_idx in range(theta_slices): + ring_points = np.zeros((phi_slices, 3)) + for phi_idx in range(phi_slices): + if (ring_idx, phi_idx) in sph_points: + if ignore == 0: + nth_min_r = np.min(np.array(sph_points[(ring_idx, phi_idx)])) # point in area with smallest radius + else: + nth_min_r = np.partition(np.array(sph_points[(ring_idx, phi_idx)]), ignore)[ignore] + else: + print("no ", (ring_idx, phi_idx)) + exit() # something wrong + midpoint_theta, midpoint_phi = grid_mid_point(ring_idx, phi_idx) # the smallest radius will be assigned to a point in the middle-ish of the area + ring_points[phi_idx] = np.array(spherical_to_cartesian([nth_min_r, midpoint_theta, midpoint_phi])) + + dayside_magnetopause[ring_idx] = ring_points + + + ### x < 0 magnetopause ### + # rest: look for magnetopause in yz-planes + ## define points in the x axis where to find magnetopause points on the yz-plane + x_points = np.linspace(0.0, end_x, x_point_n-dayside_x_point_n) + + ## interpolate more exact points for streamlines at exery x_point + new_streampoints = np.zeros((len(x_points), len(streams), 2)) # new array for keeping interpolated streamlines in form new_streampoints[x_point, streamline, y and z -coordinates] + i=0 # streamline + for stream in streams: + interpolated_streamline = interpolate(stream, x_points) + + if type(interpolated_streamline) is np.ndarray: # don't use 'discarded' streamlines, see function interpolate() + for j in range(0, len(x_points)): + y,z = interpolated_streamline[1,j], interpolated_streamline[2,j] + new_streampoints[j, i,:] = np.array([y,z]) + i += 1 + + + + ## create a list of streamline points in polar coordinates (for every x_point) + polar_coords = np.zeros_like(new_streampoints) + for i in range(0,new_streampoints.shape[0]): + for j in range(0,new_streampoints.shape[1]): + polar_coords[i,j,:] = cartesian_to_polar(new_streampoints[i,j]) + + + ## now start making the magnetopause + ## in each x_point, divide the plane into sectors and look for the closest streamline to x-axis in the sector + + sector_width = 360/sector_n + magnetopause = np.zeros((len(x_points), sector_n, 3)) + + for i,x_point in enumerate(x_points): #loop over every chosen x-axis point + # divide the yz-plane into sectors + for j, mean_sector_angle in enumerate(np.arange(0, 360, sector_width)): + min_angle = mean_sector_angle-sector_width/2 + max_angle = mean_sector_angle+sector_width/2 + + # find points that are in the sector + if mean_sector_angle == 0: # special case as the first sector needs streamlines around phi=0 + min_angle = min_angle+360 + # divide into phi<360 and phi>0 + sector1 = polar_coords[i, (polar_coords[i,:,1] <= 360)*(polar_coords[i,:,1] > min_angle)] + sector2 = polar_coords[i, (polar_coords[i,:,1] <= max_angle)*(polar_coords[i,:,1] >= 0)] + sector_points = np.concatenate((sector1, sector2)) + + else: + sector_points = polar_coords[i, (polar_coords[i,:,1] <= max_angle)*(polar_coords[i,:,1] > min_angle)] + + # discard 'points' with r=0 and check that there's at least one streamline point in the sector + sector_points = sector_points[sector_points[:,0] != 0.0] + if sector_points.size == 0: + raise ValueError('No streamlines found in the sector') + + # find the points closest to the x-axis + if ignore == 0: + nth_closest_point_radius = sector_points[sector_points[:,0].argmin(), 0] # smallest radius + else: + nth_closest_point_radius = np.partition(sector_points[:,0], ignore)[ignore] + + #closest_point_radius = np.median(sector_points[:,0]) # median radius + #if x_point < subsolar_x-2e6: + # closest_point_radius = np.median(sector_points[:,0]) # median radius + #else: + # closest_point_radius = sector_points[sector_points[:,0].argmin(), 0] # smallest radius + + # return to cartesian coordinates and save as a magnetopause point at the middle of the sector + y,z = polar_to_cartesian(nth_closest_point_radius, mean_sector_angle) + magnetopause[i,j,:] = [x_point, y, z] + + + # make a tip point for the magnetopause for prettier 3d plots + tip = np.array([subsolar_x, 0, 0]) + tips = np.tile(tip, (magnetopause.shape[1],1)) + magnetopause = np.vstack(([tips], dayside_magnetopause, magnetopause)) + + return magnetopause + + +def find_magnetopause_sw_streamline_3d(vlsvfile, streamline_seeds=None, seeds_n=25, seeds_x0=20*6371000, seeds_range=[-5*6371000, 5*6371000], dl=2e6, iterations=200, end_x=-15*6371000, x_point_n=50, sector_n=36, ignore=0): + """Finds the magnetopause position by tracing streamlines of the velocity field. + + Note: there may be a slight jump at x=0. This may be due to difference in methods (pointcloud vs. interpolation). + If the subsolar point area looks off then more streamlines near the x-axis are needed. + + + :param vlsvfile: path to .vlsv bulk file to use for VlsvReader + :kword streamline_seeds: optional streamline starting points in numpy array + :kword seeds_n: instead of streamline_seeds provide a number of streamlines to be traced + :kword seeds_x0: instead of streamline_seeds provide an x-coordinate for streamline starting points + :kword seeds_range: instead of streamline_seeds provide [min, max] range to use for streamline starting point coordinates (both y- and z-directions use the same range) + :kword dl: streamline iteration step length in m + :kword iterations: int, number of iteration steps + :kword end_x: tail end x-coordinate (how far along the x-axis the magnetopause is calculated) + :kword x_point_n: integer, how many parts the magnetopause will be divided in between the subsolar point and tail end + :kword sector_n: integer, how many sectors the magnetopause will be divided in on each yz-plane/radial sector + :kword ignore: how many inner streamlines will be ignored when calculating the magnetopause + + :returns: vertices, surface where vertices are numpy arrays in shape [[x0,y0,z0], [x1,y1,z1],...] and surface is a vtk vtkDataSetSurfaceFilter object + """ + + streams = make_streamlines(vlsvfile, streamline_seeds, seeds_n, seeds_x0, seeds_range, dl, iterations) + magnetopause = make_magnetopause(streams, end_x, x_point_n, sector_n, ignore) + vertices, faces = make_surface(magnetopause) + surface = make_vtk_surface(vertices, faces) + + return vertices, surface + + + diff --git a/analysator/pyCalculations/timeevolution.py b/analysator/pyCalculations/timeevolution.py index 3852759c..3742c5bd 100644 --- a/analysator/pyCalculations/timeevolution.py +++ b/analysator/pyCalculations/timeevolution.py @@ -192,6 +192,7 @@ def point_time_evolution( vlsvReader_list, variables, coordinates, units="", met # For optimization purposes we are now freeing vlsvReader's memory # Note: Upon reading data vlsvReader created an internal hash map that takes a lot of memory vlsvReader.optimize_clear_fileindex_for_cellid() + vlsvReader.optimize_clear_fileindex_for_cellid_blocks() # Close the vlsv reader's file: vlsvReader.optimize_close_file() from output import output_1d diff --git a/analysator/pyVlsv/vlsvcache.py b/analysator/pyVlsv/vlsvcache.py new file mode 100644 index 00000000..15f8ed98 --- /dev/null +++ b/analysator/pyVlsv/vlsvcache.py @@ -0,0 +1,215 @@ +# +# This file is part of Analysator. +# Copyright 2013-2016 Finnish Meteorological Institute +# Copyright 2017-2024 University of Helsinki +# +# For details of usage, see the COPYING file and read the "Rules of the Road" +# at http://www.physics.helsinki.fi/vlasiator/ +# +# This program is free software; you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation; either version 2 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License along +# with this program; if not, write to the Free Software Foundation, Inc., +# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. +# + +''' Utilities for caching VLSV data and metadata. +''' + +import logging +import os +import sys +import warnings +import numbers +import numpy as np +from operator import itemgetter +import re +import rtree +import time + +class VariableCache: + ''' Class for handling in-memory variable/reducer caching. + ''' + def __init__(self, reader): + self.__varcache = {} # {(varname, operator):data} + self.__reader = reader + + def keys(self): + return self.__varcache.keys() + + def __getitem__(self, key): + return self.__varcache[key] + + def __setitem__(self, key, value): + self.__varcache[key] = value + + + def read_variable_from_cache(self, name, cellids, operator): + ''' Read variable from cache instead of the vlsv file. + :param name: Name of the variable + :param cellids: a value of -1 reads all data + :param operator: Datareduction operator. "pass" does no operation on data + :returns: numpy array with the data, same format as read_variable + + .. seealso:: :func:`read_variable` + ''' + + var_data = self.__varcache[(name,operator)] + if var_data.ndim == 2: + value_len = var_data.shape[1] + else: + value_len = 1 + + if isinstance(cellids, numbers.Number): + if cellids == -1: + return var_data + else: + return var_data[self.__reader.get_cellid_locations()[cellids]] + else: + if(len(cellids) > 0): + indices = np.array(itemgetter(*cellids)(self.__reader.get_cellid_locations()),dtype=np.int64) + else: + indices = np.array([],dtype=np.int64) + if value_len == 1: + return var_data[indices] + else: + return var_data[indices,:] + +class FileCache: + ''' Top-level class for caching to file. + ''' + + def __init__(self, reader) -> None: + self.__reader = reader + + self.__rtree_index_files = [] + self.__rtree_index = None + self.__rtree_idxfile = os.path.join(self.get_cache_folder(),"rtree.idx") + self.__rtree_datfile = os.path.join(self.get_cache_folder(),"rtree.dat") + self.__rtree_properties = rtree.index.Property() + self.__rtree_properties.dimension = 3 + self.__rtree_properties.overwrite=True + + + def get_cache_folder(self): + fn = self.__reader.file_name + + head,tail = os.path.split(fn) + path = head + numslist = re.findall(r'\d+(?=\.vlsv)', tail) + + if(len(numslist) == 0): + path = os.path.join(path,"vlsvcache",tail[:-5]) + else: + nums = numslist[-1] + head, tail = tail.split(nums) + + leading_zero = True + path = os.path.join(path,"vlsvcache",head[:-1]) + for i,n in enumerate(nums): + if n == '0' and leading_zero and i!=len(nums)-1: + continue + if leading_zero: + fmt = "{:07d}" + else: + fmt = "{:0"+str(7-i)+"d}" + path = os.path.join(path, fmt.format(int(n)*10**(len(nums)-i-1))) + leading_zero = False + + return path + + def clear_cache_folder(self): + path = self.get_cache_folder() + import shutil + shutil.rmtree(path) + + def set_cellid_spatial_index(self, force = False): + if not os.path.exists(self.get_cache_folder()): + os.makedirs(self.get_cache_folder()) + + if force: + if os.path.exists(self.__rtree_idxfile): + os.remove(self.__rtree_idxfile) + if os.path.exists(self.__rtree_datfile): + os.remove(self.__rtree_datfile) + if(not os.path.isfile(self.__rtree_idxfile) or not os.path.isfile(self.__rtree_datfile)): + t0 = time.time() + + bboxes = self.__reader.get_mesh_domain_extents("SpatialGrid") + bboxes = bboxes.reshape((-1,6), order='C') + + self.__rtree_index = rtree.index.Index(self.__rtree_idxfile[:-4],properties=self.__rtree_properties, interleaved=False) + for rank, bbox in enumerate(bboxes): + self.__rtree_index.insert(rank, bbox) + + logging.info("index set in "+str(time.time()-t0)+" seconds") + else: + pass + + def get_cellid_spatial_index(self, force = False): + if self.__rtree_index == None: + if(force): + self.set_cellid_spatial_index(force) + elif not os.path.isfile(self.__rtree_idxfile) or not os.path.isfile(self.__rtree_datfile): + self.__rtree_index = None + else: + self.__rtree_index = rtree.index.Index(self.__rtree_idxfile[:-4], properties=self.__rtree_properties, interleaved=False) + + return self.__rtree_index + + +class MetadataFileCache(FileCache): + ''' File caching class for storing "lightweight" metadata. + ''' + pass + # superclass constructor called instead if no __init__ here + # def __init__(self, reader) -> None: + # super(MetadataFileCache, self).__init__(reader) + + + +class VariableFileCache(FileCache): + ''' File caching class for storing intermediate data, such as + gradient terms that are more expensive to compute, over whole grids with + more HDD footprint. + ''' + pass + # superclass constructor called instead if no __init__ here + # def __init__(self, reader) -> None: + # super(MetadataFileCache, self).__init__(reader) + +class PicklableFile(object): + ''' Picklable file pointer object. + ''' + def __init__(self, fileobj): + self.fileobj = fileobj + + def __getattr__(self, key): + return getattr(self.fileobj, key) + + def __getstate__(self): + ret = self.__dict__.copy() + ret['_file_name'] = self.fileobj.name + ret['_file_mode'] = self.fileobj.mode + if self.fileobj.closed: + ret['_file_pos'] = 0 + else: + ret['_file_pos'] = self.fileobj.tell() + del ret['fileobj'] + return ret + + def __setstate__(self, dict): + self.fileobj = open(dict['_file_name'], dict['_file_mode']) + self.fileobj.seek(dict['_file_pos']) + del dict['_file_name'] + del dict['_file_mode'] + del dict['_file_pos'] + self.__dict__.update(dict) diff --git a/analysator/pyVlsv/vlsvreader.py b/analysator/pyVlsv/vlsvreader.py index ac0caf88..780087e3 100644 --- a/analysator/pyVlsv/vlsvreader.py +++ b/analysator/pyVlsv/vlsvreader.py @@ -1,4 +1,4 @@ -#s +# # This file is part of Analysator. # Copyright 2013-2016 Finnish Meteorological Institute # Copyright 2017-2024 University of Helsinki @@ -30,8 +30,11 @@ import sys import re import numbers +import pickle # for caching linked readers, switch to VLSV/XML at some point - h5py? +# import h5py import vlsvvariables +import vlsvcache from reduction import datareducers,multipopdatareducers,data_operators,v5reducers,multipopv5reducers,deprecated_datareducers try: from collections.abc import Iterable @@ -48,31 +51,7 @@ interp_method_aliases = {"trilinear":"linear"} -class PicklableFile(object): - def __init__(self, fileobj): - self.fileobj = fileobj - - def __getattr__(self, key): - return getattr(self.fileobj, key) - - def __getstate__(self): - ret = self.__dict__.copy() - ret['_file_name'] = self.fileobj.name - ret['_file_mode'] = self.fileobj.mode - if self.fileobj.closed: - ret['_file_pos'] = 0 - else: - ret['_file_pos'] = self.fileobj.tell() - del ret['fileobj'] - return ret - - def __setstate__(self, dict): - self.fileobj = open(dict['_file_name'], dict['_file_mode']) - self.fileobj.seek(dict['_file_pos']) - del dict['_file_name'] - del dict['_file_mode'] - del dict['_file_pos'] - self.__dict__.update(dict) +neighbors_cache_file = "neighbors_cache.pkl" def dict_keys_exist(dictionary, query_keys, prune_unique=False): if query_keys.shape[0] == 0: @@ -114,7 +93,7 @@ def fsGlobalIdToGlobalIndex(globalids, bbox): # Read in the global ids and indices for FsGrid cells, returns # min and max corners of the fsGrid chunk by rank def fsReadGlobalIdsPerRank(reader): - numWritingRanks = reader.read_parameter("numWritingRanks") + numWritingRanks = reader.get_numWritingRanks("fsgrid") rawData = reader.read(tag="MESH", name="fsgrid") bbox = reader.read(tag="MESH_BBOX", mesh="fsgrid") sizes = reader.read(tag="MESH_DOMAIN_SIZES", mesh="fsgrid") @@ -181,27 +160,44 @@ def __del__(self): if (hasattr(self, "__fptr")) and self.__fptr is not None: self.__fptr.close() - def __init__(self, file_name, fsGridDecomposition=None): + def __init__(self, file_name, fsGridDecomposition=None, file_cache = 0): ''' Initializes the vlsv file (opens the file, reads the file footer and reads in some parameters) :param file_name: Name of the vlsv file - :param fsGridDecomposition: Either None or a len-3 list of ints. + :kword fsGridDecomposition: Either None or a len-3 list of ints [None]. List (length 3): Use this as the decomposition directly. Product needs to match numWritingRanks. + :kword file_cache: Boolean, [False]: cache slow-to-compute data to disk (:seealso get_cache_folder) ''' # Make sure the path is set in file name: file_name = os.path.abspath(file_name) self.file_name = file_name + self.file_cache = file_cache try: - self.__fptr = PicklableFile(open(self.file_name,"rb")) + self.__fptr = vlsvcache.PicklableFile(open(self.file_name,"rb")) except FileNotFoundError as e: logging.info("File not found: " + self.file_name) raise e self.__xml_root = ET.fromstring("") self.__fileindex_for_cellid={} + self.__full_fileindex_for_cellid = False + self.__cellid_spatial_index=None + self.__rankwise_fileindex_for_cellid = {} # { : {cellid: offset}} + self.__loaded_fileindex_ranks = set() + + self.__metadata_cache = vlsvcache.MetadataFileCache(self) + self.__metadata_read = False + self.__metadata_dict = {} # Used for reading in and storing derived/other metadata such as linked file paths + self.__linked_files = set() + self.__linked_readers = set() + + + + self.__mesh_domain_sizes = {} self.__max_spatial_amr_level = -1 + self.__grid_epsilon = None self.__fsGridDecomposition = fsGridDecomposition self.use_dict_for_blocks = False @@ -212,7 +208,8 @@ def __init__(self, file_name, fsGridDecomposition=None): self.__order_for_cellid_blocks = {} # per-pop self.__vg_indexes_on_fg = np.array([]) # SEE: map_vg_onto_fg(self) - self.variable_cache = {} # {(varname, operator):data} + self.__variable_cache = vlsvcache.VariableCache(self) # {(varname, operator):data} + self.__params_cache = {} # {name:data} self.__pops_init = False @@ -220,7 +217,7 @@ def __init__(self, file_name, fsGridDecomposition=None): self.__unavailable_reducers = set() # Set of strings of datareducer names self.__current_reducer_tree_nodes = set() # Set of strings of datareducer names - self.__read_xml_footer() + # vertex-indices is a 3-tuple of integers self.__dual_cells = {(0,0,0):(1,1,1,1,1,1,1,1)} # vertex-indices : 8-tuple of cellids at each corner (for x for y for z) self.__dual_bboxes = {} # vertex-indices : 6-list of (xmin, ymin, zmin, xmax, ymax, zmax) for the bounding box of each dual cell @@ -229,6 +226,13 @@ def __init__(self, file_name, fsGridDecomposition=None): self.__cell_neighbours = {} # cellid : set of cellids (all neighbors sharing a vertex) self.__cell_duals = {} # cellid : tuple of vertex-indices that span this cell self.__regular_neighbor_cache = {} # cellid-of-low-corner : (8,) np.array of cellids) + self.__neighbors_cache_len = 0 + self.__neighbors_cache_available = os.path.isfile(os.path.join(self.get_cache_folder(),neighbors_cache_file)) + self.__neighbors_cache_loaded = False + + # Start calling functions only after initializing trivial members + self.get_linked_readers() + self.__read_xml_footer() # Check if the file is using new or old vlsv format # Read parameters (Note: Reading the spatial cell locations and @@ -315,6 +319,150 @@ def __init__(self, file_name, fsGridDecomposition=None): self.__fptr.close() + def get_grid_epsilon(self): + if self.__grid_epsilon is None: + # one-thousandth of the max refined cell; self.get_max_refinement_level() however reads all cellids, so just temp here by assuming 8 refinement levels.. which is plenty for now + self.__grid_epsilon = 1e-3*np.array([self.__dx, self.__dy, self.__dz])/2**8 + return self.__grid_epsilon + + + def get_metadata_filename(self): + pth, base = os.path.split(self.file_name) + + s = os.path.join(pth,"vlsvmeta",base[:-5]+"_metadata.pkl") + return s + + # def get_h5_metadata(self, key, default): + # ''' Read metadata from hdf5 metadata file, and if not available, + # return the given default value. + + # :param data: str, a key to stored metadata. + # :param default: value to return if key does not exist + + # ''' + + # if type(key) == type(("a tuple",)): + # print("tuple reader not implemented") + # elif type(key) == type("a string"): + # print("Reading str-keyed data") + # else: + # raise TypeError("key must be str or tuple") + + # if not self.__metadata_read: + # try: + # fn = self.get_metadata_filename() + # with open(fn,'rb') as f: + # self.__metadata_dict = pickle.load(f) + # except: + # logging.debug("No metadata file found.") + # self.__metadata_read = True + + # return self.__metadata_dict.get(key,default) + + + def get_reader_metadata(self, key, default): + ''' Read metadata from metadata file, and if not available, + return the given default value. + + :param data: str, a key to stored metadata. + :param default: value to return if key does not exist + + ''' + + + if not self.__metadata_read: + try: + fn = self.get_metadata_filename() + with open(fn,'rb') as f: + self.__metadata_dict = pickle.load(f) + except: + logging.debug("No metadata file found.") + self.__metadata_read = True + + return self.__metadata_dict.get(key,default) + + def add_metadata(self, key, value): + self.__metadata_dict[key] = value + self.save_metadata() + + def save_metadata(self): + fn = self.get_metadata_filename() + try: + with open(fn,'wb') as f: + pickle.dump(self.__metadata_dict,f) + except Exception as e: + logging.warning("Could not save metadata file, error: "+str(e)) + + def get_linked_readers_filename(self): + '''Need to go to a consolidated metadata handler''' + pth, base = os.path.split(self.file_name) + + s = os.path.join(self.__metadata_cache.get_cache_folder(),base[:-5]+"_linked_readers.txt") + return s + + def get_linked_readers(self, reload=False): + self.__linked_files = self.get_reader_metadata("linked_reader_files", set()) + if len(self.__linked_files)==0 or reload: + if(os.path.isfile(self.get_linked_readers_filename())): + with open(self.get_linked_readers_filename(), 'r') as f: + l = f.readlines() + logging.info("Loaded linked readers from "+self.get_linked_readers_filename()) + self.__linked_files.update(l) + print(l) + + else: + self.add_linked_readers() + + self.add_metadata("linked_reader_files",self.__linked_files) + + return self.__linked_readers + + + def add_linked_file(self, fname): + if os.path.exists(fname): + self.__linked_files.add(VlsvReader(fname)) + else: + logging.warning("Could not link "+fname+" (path does not exist)") + + def add_linked_reader(self, fname): + if os.path.exists(fname): + for reader in self.__linked_readers: + if fname == reader.file_name: + return + self.__linked_readers.add(VlsvReader(fname)) + else: + logging.warning("Could not link "+fname+" (path does not exist)") + + def add_linked_readers(self): + for fname in self.__linked_files: + self.add_linked_reader(fname) + + def save_linked_readers_file(self, overwrite = False): + fn = self.get_linked_readers_filename() + if not overwrite: # Load existing linked reader to not overwrite everything + self.get_linked_readers() + + logging.info("Saving linked readers to "+fn) + dn = os.path.dirname(fn) + if not os.path.isdir(dn): + os.mkdir(dn) + with open(fn,'w') as f: + lines = [] + for line in self.__linked_files: + lines += (line+os.linesep) # gather from set and insert line swap + f.writelines(lines) + + def clear_linked_readers(self): + self.__linked_files.clear() + self.__linked_readers.clear() + + def clear_linked_readers_file(self): + fn = self.get_linked_readers_filename() + if os.path.exists(fn): + os.remove(fn) + else: + logging.info("Tried to remove nonexisting linked reader file "+fn) + def __popmesh(self, popname): ''' Get the population-specific vspace mesh info object, and initialize it if it does not exist @@ -470,9 +618,10 @@ def __read_xml_footer(self): def __read_fileindex_for_cellid(self): """ Read in the cell ids and create an internal dictionary to give the index of an arbitrary cellID """ - if not self.__fileindex_for_cellid == {}: + if self.__full_fileindex_for_cellid: return - + + # print("fileindex!") cellids=self.read(mesh="SpatialGrid",name="CellID", tag="VARIABLE") #Check if it is not iterable. If it is a scale then make it a list @@ -481,6 +630,7 @@ def __read_fileindex_for_cellid(self): # self.__fileindex_for_cellid = {cellid:index for index,cellid in enumerate(cellids)} for index,cellid in enumerate(cellids): self.__fileindex_for_cellid[cellid] = index + self.__full_fileindex_for_cellid = True def __read_blocks(self, cellid, pop="proton"): ''' Read raw velocity block data from the open file. @@ -691,18 +841,24 @@ def __check_datareducer(self, name, reducer): def get_variables(self): - varlist = [] + varlist = set() + + for reader in self.__linked_readers: + varlist.update(set(reader.get_variables())) for child in self.__xml_root: if child.tag == "VARIABLE" and "name" in child.attrib: name = child.attrib["name"] - varlist.append(name) + varlist.add(name) - return varlist + return list(varlist) def get_reducers(self): - varlist = [] + varlist = set() + + for reader in self.__linked_readers: + varlist.update(reader.get_reducers()) reducer_max_len = 0 @@ -718,11 +874,11 @@ def get_reducers(self): if self.__check_datareducer(name,reducer): if name[:3] == 'pop': for pop in self.active_populations: - varlist.append(pop+'/'+name[4:]) + varlist.add(pop+'/'+name[4:]) else: - varlist.append(name) + varlist.add(name) - return varlist + return list(varlist) def list(self, parameter=True, variable=True, mesh=False, datareducer=False, operator=False, other=False): @@ -807,6 +963,27 @@ def check_parameter( self, name ): return True return False + def linked_readers_check_variable(self, name): + ''' Test all linked variables if any of them returns True on test function + + :param fun: Function to pass to linked readers (VlsvReader member function/first param VlsvReader) + :param args: arguments to pass to fun + :param kwargs: keyword arguments to pass to fun + ''' + + ret = False + for reader in self.__linked_readers: + try: + ret = reader.check_variable(name) + except Exception as e: # Let's not care if a sidecar file does not contain something + pass + if ret: + return ret + else: + continue + + return False + def check_variable( self, name ): ''' Checks if a given variable is in the vlsv reader @@ -826,6 +1003,9 @@ def check_variable( self, name ): # Variable not in the vlsv file plot_B_vol() ''' + if self.linked_readers_check_variable(name): + return True + for child in self.__xml_root: if child.tag == "VARIABLE" and "name" in child.attrib: if child.attrib["name"].lower() == name.lower(): @@ -888,6 +1068,68 @@ def get_cellid_locations(self): self.__read_fileindex_for_cellid() return self.__fileindex_for_cellid + def get_mesh_domain_sizes(self, mesh): + + if mesh not in self.__mesh_domain_sizes.keys(): + self.__mesh_domain_sizes[mesh] = self.read(name="", tag="MESH_DOMAIN_SIZES", mesh=mesh) + return self.__mesh_domain_sizes[mesh] + + def get_numWritingRanks(self, mesh): + ''' Get the number of writing ranks. This does not exist for all outputfiles eerywhere. + ''' + ranks = None + if "numWritingRanks" in self.__params_cache.keys(): + return self.__params_cache["numWritingRanks"] + else: + try: + ranks = self.read_parameter("numWritingRanks") + except: + ranks = self.read_attribute(mesh=mesh, attribute="arraysize", tag="MESH_DOMAIN_SIZES") + self.__params_cache["numWritingRanks"] = ranks + + return ranks + + + def __read_fileindex_for_cellid_rank(self, rank): + """ Read in the cell ids and create an internal dictionary to give the index of an arbitrary cellID + """ + # print("filling rank",rank) + + if rank > self.get_numWritingRanks("SpatialGrid"): + raise ValueError("Tried to read rank "+rank+" out of "+self.get_numWritingRanks("SpatialGrid")) + + if not self.__rankwise_fileindex_for_cellid.get(rank,{}) == {}: + return + + mesh_domain_sizes = self.get_mesh_domain_sizes("SpatialGrid") + n_domain_cells = mesh_domain_sizes[:,0]-mesh_domain_sizes[:,1] + domain_offsets = np.cumsum(np.hstack((np.array([0]),n_domain_cells[:-1]))) + + cellids = self.read_cellids_with_offset(domain_offsets[rank],n_domain_cells[rank]) + + #Check if it is not iterable. If it is a scale then make it a list + if(not isinstance(cellids, Iterable)): + cellids=[ cellids ] + self.__rankwise_fileindex_for_cellid[rank] = {cellid: domain_offsets[rank]+index for index,cellid in enumerate(cellids)} + self.__loaded_fileindex_ranks.add(rank) + self.__fileindex_for_cellid.update(self.__rankwise_fileindex_for_cellid[rank]) + + + def get_cellid_locations_rank(self, rank): + ''' Returns a dictionary with cell id as the key and the index of the cell id as the value. + So far this pipeline will update the main fileindex, as it is not too heavy an operation - but + moving towards handling data more rankwise may be nice some time in the future. + + :param rank: Find (and initialize) fileindex for cells contained by this rank + + ''' + # if len( self.__fileindex_for_cellid ) == 0: + self.__read_fileindex_for_cellid_rank(rank) + return self.__rankwise_fileindex_for_cellid[rank] + + + + def print_version(self): ''' Prints version information from VLSV file. @@ -1057,6 +1299,105 @@ def read_attribute(self, name="", mesh="", attribute="", tag=""): raise ValueError("Variable or attribute not found") + def read_with_offset(self, datatype,variable_offset, read_size, read_offsets, element_size, vector_size): + if self.__fptr.closed: + fptr = open(self.file_name,"rb") + else: + fptr = self.__fptr + if len(read_offsets) !=1: + arraydata = [] + for r_offset in read_offsets: + use_offset = int(variable_offset + r_offset) + fptr.seek(use_offset) + if datatype == "float" and element_size == 4: + data = np.fromfile(fptr, dtype = np.float32, count=vector_size*read_size) + if datatype == "float" and element_size == 8: + data = np.fromfile(fptr, dtype=np.float64, count=vector_size*read_size) + if datatype == "int" and element_size == 4: + data = np.fromfile(fptr, dtype=np.int32, count=vector_size*read_size) + if datatype == "int" and element_size == 8: + data = np.fromfile(fptr, dtype=np.int64, count=vector_size*read_size) + if datatype == "uint" and element_size == 4: + data = np.fromfile(fptr, dtype=np.uint32, count=vector_size*read_size) + if datatype == "uint" and element_size == 8: + data = np.fromfile(fptr, dtype=np.uint64, count=vector_size*read_size) + if len(read_offsets)!=1: + arraydata.append(data) + if len(read_offsets) !=1: + data = np.array(arraydata) + + return data + + def get_read_offsets(self, cellids, array_size, element_size, vector_size): + ''' Figure out somewhat optimal offsets for read operations + ''' + + # if isinstance(cellids, numbers.Number): + # if cellids >= 0: + # read_filelayout = True + # else: # cellids = -1 + # read_filelayout = False + # else: + # read_filelayout = True + # # Here could be a conditional optimization for unique CellIDs, + # # but it actually makes a very large query slower. + + # if (len( self.__fileindex_for_cellid ) == 0) and read_filelayout: + # # Do we need to construct the cellid index? + # self.__read_fileindex_for_cellid() + # Define efficient method to read data in + reorder_data = False + if not isinstance(cellids, numbers.Number): + # Read multiple specified cells + # If we're reading a large amount of single cells, it'll be faster to just read all + # data from the file system and sort through it. For the CSC disk system, this + # becomes more efficient for over ca. 5000 cellids. + if len(cellids) >5000: # TODO re-evaluate threshold + reorder_data = True + result_size = len(cellids) + read_size = array_size + read_offsets = [0] + else: # Read multiple cell ids one-by-one + try: + result_size = len(cellids) + read_size = 1 + read_offsets = [self.__fileindex_for_cellid[cid]*element_size*vector_size for cid in cellids] + except: + self.__read_fileindex_for_cellid() + result_size = len(cellids) + read_size = 1 + read_offsets = [self.__fileindex_for_cellid[cid]*element_size*vector_size for cid in cellids] + else: # single cell or all cells + if cellids < 0: # -1, read all cells + result_size = array_size + read_size = array_size + read_offsets = [0] + else: # single cell id + if self.__full_fileindex_for_cellid: # Fileindex already exists, we might as well use it + result_size = 1 + read_size = 1 + read_offsets = [self.__fileindex_for_cellid[cellids]*element_size*vector_size] + elif self.__cellid_spatial_index != None: + # Here a faster alternative + result_size = 1 + read_size = 1 + eps = self.get_grid_epsilon() + qp = self.get_cell_coordinates(cellids) + qqp = np.hstack((qp-eps,qp+eps)) + pq = np.array([qqp[0],qqp[3],qqp[1],qqp[4],qqp[2],qqp[5]]) + rankids = self.__cellid_spatial_index.intersection() + read_offsets = None + for rankid in rankids: + indexdict = self.get_cellid_locations_rank(rankid) + if cellids in indexdict.keys(): + read_offsets = [indexdict[cellids]*element_size*vector_size] + else: + self.__read_fileindex_for_cellid() + result_size = 1 + read_size = 1 + read_offsets = [self.__fileindex_for_cellid[cellids]*element_size*vector_size] + + return read_size, read_offsets, result_size, reorder_data def read(self, name="", tag="", mesh="", operator="pass", cellids=-1): ''' Read data from the open vlsv file. @@ -1080,24 +1421,9 @@ def read(self, name="", tag="", mesh="", operator="pass", cellids=-1): # Force lowercase name for internal checks name = name.lower() - if tag == "VARIABLE": - if (name,operator) in self.variable_cache.keys(): + if (name,operator) in self.__variable_cache.keys(): return self.read_variable_from_cache(name, cellids, operator) - - if isinstance(cellids, numbers.Number): - if cellids >= 0: - read_filelayout = True - else: # cellids = -1 - read_filelayout = False - else: - read_filelayout = True - # Here could be a conditional optimization for unique CellIDs, - # but it actually makes a very large query slower. - - if (len( self.__fileindex_for_cellid ) == 0) and read_filelayout: - # Do we need to construct the cellid index? - self.__read_fileindex_for_cellid() # Get population and variable names from data array name if '/' in name: @@ -1136,58 +1462,15 @@ def read(self, name="", tag="", mesh="", operator="pass", cellids=-1): datatype = child.attrib["datatype"] variable_offset = ast.literal_eval(child.text) - # Define efficient method to read data in - reorder_data = False - arraydata = [] - try: # try-except to see how many cellids were given - lencellids=len(cellids) - # Read multiple specified cells - # If we're reading a large amount of single cells, it'll be faster to just read all - # data from the file system and sort through it. For the CSC disk system, this - # becomes more efficient for over ca. 5000 cellids. - if lencellids>5000: - reorder_data = True - result_size = len(cellids) - read_size = array_size - read_offsets = [0] - else: # Read multiple cell ids one-by-one - result_size = len(cellids) - read_size = 1 - read_offsets = [self.__fileindex_for_cellid[cid]*element_size*vector_size for cid in cellids] - except: # single cell or all cells - if cellids < 0: # -1, read all cells - result_size = array_size - read_size = array_size - read_offsets = [0] - else: # single cell id - result_size = 1 - read_size = 1 - read_offsets = [self.__fileindex_for_cellid[cellids]*element_size*vector_size] - - if self.__fptr.closed: - fptr = open(self.file_name,"rb") + if not isinstance(cellids, numbers.Number): + cellids = np.array(cellids, dtype=np.int64) + cellids_nonzero = cellids[cellids!=0] else: - fptr = self.__fptr + cellids_nonzero = cellids - for r_offset in read_offsets: - use_offset = int(variable_offset + r_offset) - fptr.seek(use_offset) - if datatype == "float" and element_size == 4: - data = np.fromfile(fptr, dtype = np.float32, count=vector_size*read_size) - if datatype == "float" and element_size == 8: - data = np.fromfile(fptr, dtype=np.float64, count=vector_size*read_size) - if datatype == "int" and element_size == 4: - data = np.fromfile(fptr, dtype=np.int32, count=vector_size*read_size) - if datatype == "int" and element_size == 8: - data = np.fromfile(fptr, dtype=np.int64, count=vector_size*read_size) - if datatype == "uint" and element_size == 4: - data = np.fromfile(fptr, dtype=np.uint32, count=vector_size*read_size) - if datatype == "uint" and element_size == 8: - data = np.fromfile(fptr, dtype=np.uint64, count=vector_size*read_size) - if len(read_offsets)!=1: - arraydata.append(data) - - fptr.close() + read_size, read_offsets, result_size, reorder_data = self.get_read_offsets(cellids_nonzero, array_size, element_size, vector_size) + + data = self.read_with_offset(datatype, variable_offset, read_size, read_offsets, element_size, vector_size) if len(read_offsets)==1 and reorder_data: # Many single cell id's requested @@ -1195,17 +1478,26 @@ def read(self, name="", tag="", mesh="", operator="pass", cellids=-1): arraydata = np.array(data) if vector_size > 1: arraydata=arraydata.reshape(-1, vector_size) - - append_offsets = [self.__fileindex_for_cellid[cid] for cid in cellids] + + mask = ~dict_keys_exist(self.__fileindex_for_cellid, cellids_nonzero) + + self.do_partial_fileindex_update(self.get_cell_coordinates(cellids_nonzero[mask])) + + append_offsets = [self.__fileindex_for_cellid[cid] for cid in cellids_nonzero] data = arraydata[append_offsets,...] data = np.squeeze(data) if len(read_offsets)!=1: # Not-so-many single cell id's requested - data = np.squeeze(np.array(arraydata)) + data = np.squeeze(np.array(data)) if vector_size > 1: data=data.reshape(result_size, vector_size) + + if not isinstance(cellids, numbers.Number): + data_out = np.full_like(data, np.nan, shape=(len(cellids),*data.shape[1:])) + data_out[cellids!=0,...] = data + data = data_out # If variable vector size is 1, and requested magnitude, change it to "absolute" if vector_size == 1 and operator=="magnitude": @@ -1325,9 +1617,49 @@ def read(self, name="", tag="", mesh="", operator="pass", cellids=-1): return data_operators[operator](reducer.operation( tmp_vars )) if name!="": - raise ValueError("Error: variable "+name+"/"+tag+"/"+mesh+"/"+operator+" not found in .vlsv file or in data reducers!") + raise ValueError("Error: variable "+name+"/"+tag+"/"+mesh+"/"+operator+" not found in .vlsv file or in data reducers!\n Reader file "+self.file_name) + + def read_cellids_with_offset(self, start, ncells): + import ast + tag="VARIABLE" + name="cellid" + mesh="SpatialGrid" + if self.__fptr.closed: + fptr = open(self.file_name,"rb") + else: + fptr = self.__fptr + # Seek for requested data in VLSV file + for child in self.__xml_root: + if tag != "": + if child.tag != tag: + continue + # Verify that any requested name or mesh matches those of the data + if name != "": + if not "name" in child.attrib: + continue + if child.attrib["name"].lower() != name: + continue + if mesh != "": + if not "mesh" in child.attrib: + continue + if child.attrib["mesh"] != mesh: + continue + if child.tag == tag: + # Found the requested data entry in the file + vector_size = ast.literal_eval(child.attrib["vectorsize"]) + array_size = ast.literal_eval(child.attrib["arraysize"]) + element_size = ast.literal_eval(child.attrib["datasize"]) + datatype = child.attrib["datatype"] + variable_offset = ast.literal_eval(child.text) + + result_size = ncells + read_size = ncells + read_offsets = [start*element_size*vector_size] + + data = self.read_with_offset(datatype, variable_offset, read_size, read_offsets, element_size, vector_size) + return data def read_metadata(self, name="", tag="", mesh=""): ''' Read variable metadata from the open vlsv file. @@ -1595,7 +1927,10 @@ def read_interpolated_variable(self, name, coords, operator="pass",periodic=[Tru closest_cell_ids = self.get_cellid(coordinates) if method.lower() == "nearest": - final_values = self.read_variable(name, cellids=closest_cell_ids, operator=operator) + if(coordinates.shape[0] > 1): + final_values = self.read_variable(name, cellids=closest_cell_ids, operator=operator) + else: # no need to re-read if we did the test_variable already! + final_values = test_variable if stack: return final_values.squeeze() else: @@ -1801,6 +2136,37 @@ def read_fsgrid_variable_cellid(self, name, cellids=-1, operator="pass"): else: return [self.downsample_fsgrid_subarray(cid, var) for cid in cellids] + def get_fsgrid_decomposition(self): + # Try if in metadata + self.__fsGridDecomposition = self.get_reader_metadata(("MESH_DECOMPOSITION","fsgrid"),None) + if(self.__fsGridDecomposition is not None): + print("read ",self.__fsGridDecomposition) + + if self.__fsGridDecomposition is None: + self.__fsGridDecomposition = self.read(tag="MESH_DECOMPOSITION",mesh='fsgrid') + if self.__fsGridDecomposition is not None: + logging.info("Found FsGrid decomposition from vlsv file: " + str(self.__fsGridDecomposition)) + else: + logging.info("Did not find FsGrid decomposition from vlsv file.") + + # If decomposition is None even after reading, we need to calculate it: + if self.__fsGridDecomposition is None: + logging.info("Calculating fsGrid decomposition from the file") + self.__fsGridDecomposition = fsDecompositionFromGlobalIds(self) + logging.info("Computed FsGrid decomposition to be: " + str(self.__fsGridDecomposition)) + else: + # Decomposition is a list (or fail assertions below) - use it instead + pass + + numWritingRanks = self.get_numWritingRanks("SpatialGrid") + assert len(self.__fsGridDecomposition) == 3, "Manual FSGRID decomposition should have three elements, but is "+str(self.__fsGridDecomposition) + assert np.prod(self.__fsGridDecomposition) == numWritingRanks, "Manual FSGRID decomposition should have a product of numWritingRanks ("+str(numWritingRanks)+"), but is " + str(np.prod(self.__fsGridDecomposition)) + " for decomposition "+str(self.__fsGridDecomposition) + + self.add_metadata(("MESH_DECOMPOSITION","fsgrid"), self.__fsGridDecomposition) + + return self.__fsGridDecomposition + + def read_fsgrid_variable(self, name, operator="pass"): ''' Reads fsgrid variables from the open vlsv file. Arguments: @@ -1819,7 +2185,7 @@ def read_fsgrid_variable(self, name, operator="pass"): rawData = self.read(mesh='fsgrid', name=name, tag="VARIABLE", operator=operator) # Determine fsgrid domain decomposition - numWritingRanks = self.read_parameter("numWritingRanks") + numWritingRanks = self.get_numWritingRanks("SpatialGrid") if len(rawData.shape) > 1: orderedData = np.zeros([bbox[0],bbox[1],bbox[2],rawData.shape[1]]) else: @@ -1841,26 +2207,7 @@ def calcLocalSize(globalCells, ntasks, my_n): return n_per_task currentOffset = 0; - if self.__fsGridDecomposition is None: - self.__fsGridDecomposition = self.read(tag="MESH_DECOMPOSITION",mesh='fsgrid') - if self.__fsGridDecomposition is not None: - logging.info("Found FsGrid decomposition from vlsv file: " + str(self.__fsGridDecomposition)) - else: - logging.info("Did not find FsGrid decomposition from vlsv file.") - - # If decomposition is None even after reading, we need to calculate it: - if self.__fsGridDecomposition is None: - logging.info("Calculating fsGrid decomposition from the file") - self.__fsGridDecomposition = fsDecompositionFromGlobalIds(self) - logging.info("Computed FsGrid decomposition to be: " + str(self.__fsGridDecomposition)) - else: - # Decomposition is a list (or fail assertions below) - use it instead - pass - - assert len(self.__fsGridDecomposition) == 3, "Manual FSGRID decomposition should have three elements, but is "+str(self.__fsGridDecomposition) - assert np.prod(self.__fsGridDecomposition) == numWritingRanks, "Manual FSGRID decomposition should have a product of numWritingRanks ("+str(numWritingRanks)+"), but is " + str(np.prod(self.__fsGridDecomposition)) + " for decomposition "+str(self.__fsGridDecomposition) - - + self.get_fsgrid_decomposition() for i in range(0,numWritingRanks): x = (i // self.__fsGridDecomposition[2]) // self.__fsGridDecomposition[1] @@ -1961,26 +2308,7 @@ def read_variable_from_cache(self, name, cellids, operator): .. seealso:: :func:`read_variable` ''' - var_data = self.variable_cache[(name,operator)] - if var_data.ndim == 2: - value_len = var_data.shape[1] - else: - value_len = 1 - - if isinstance(cellids, numbers.Number): - if cellids == -1: - return var_data - else: - return var_data[self.__fileindex_for_cellid[cellids]] - else: - if(len(cellids) > 0): - indices = np.array(itemgetter(*cellids)(self.__fileindex_for_cellid),dtype=np.int64) - else: - indices = np.array([],dtype=np.int64) - if value_len == 1: - return var_data[indices] - else: - return var_data[indices,:] + return self.__variable_cache.read_variable_from_cache(name,cellids,operator) def read_variable_to_cache(self, name, operator="pass"): @@ -1992,9 +2320,10 @@ def read_variable_to_cache(self, name, operator="pass"): ''' # add data to dict, use a tuple of (name,operator) as the key [tuples are immutable and hashable] - self.variable_cache[(name,operator)] = self.read_variable(name, cellids=-1,operator=operator) + self.__variable_cache[(name,operator)] = self.read_variable(name, cellids=-1,operator=operator) # Also initialize the fileindex dict at the same go because it is very likely something you want to have for accessing cached values self.__read_fileindex_for_cellid() + return self.__variable_cache[(name,operator)] def read_variable(self, name, cellids=-1,operator="pass"): ''' Read variables from the open vlsv file. @@ -2007,6 +2336,8 @@ def read_variable(self, name, cellids=-1,operator="pass"): .. seealso:: :func:`read` :func:`read_variable_info` ''' cellids = get_data(cellids) + if((name,operator) in self.__variable_cache.keys()): + return self.__variable_cache.read_variable_from_cache(name,cellids,operator) # Wrapper, check if requesting an fsgrid variable if (self.check_variable(name) and (name.lower()[0:3]=="fg_")): @@ -2022,8 +2353,14 @@ def read_variable(self, name, cellids=-1,operator="pass"): return False return self.read_ionosphere_variable(name=name, operator=operator) - if((name,operator) in self.variable_cache.keys()): - return self.read_variable_from_cache(name,cellids,operator) + + for reader in self.__linked_readers: + try: + res = reader.read_variable(name=name, cellids=cellids, operator=operator) + print(self.file_name, 'read_variable', name, res) + return res + except: + pass # Passes the list of cell id's onwards - optimization for reading is done in the lower level read() method return self.read(mesh="SpatialGrid", name=name, tag="VARIABLE", operator=operator, cellids=cellids) @@ -2117,16 +2454,20 @@ def get_max_refinement_level(self): ''' Returns the maximum refinement level of the AMR ''' if self.__max_spatial_amr_level < 0: - # Read the file index for cellid - cellids=self.read(mesh="SpatialGrid",name="CellID", tag="VARIABLE") - maxcellid = np.int64(np.amax([cellids])) - - AMR_count = np.int64(0) - while (maxcellid > 0): - maxcellid -= 2**(3*(AMR_count))*(self.__xcells*self.__ycells*self.__zcells) - AMR_count += 1 - - self.__max_spatial_amr_level = AMR_count - 1 + try: + self.__max_spatial_amr_level = self.read_attribute(name="SpatialGrid", attribute="max_refinement_level",tag="MESH") + except: + # Read the file index for cellid - should maybe rather be runtime? + cellids=self.read(mesh="SpatialGrid",name="CellID", tag="VARIABLE") + maxcellid = np.int64(np.amax([cellids])) + + AMR_count = np.int64(0) + while (maxcellid > 0): + maxcellid -= 2**(3*(AMR_count))*(self.__xcells*self.__ycells*self.__zcells) + AMR_count += 1 + + self.__max_spatial_amr_level = AMR_count - 1 + return self.__max_spatial_amr_level def get_amr_level(self,cellid): @@ -2387,6 +2728,73 @@ def get_unique_cellids(self, coords): cidsout = np.array(list(OrderedDict.fromkeys(cids))) return cidsout + def do_partial_fileindex_update(self, coords): + ''' Ensure the cellids corresponding to coords or within query_window are mapped in the __fileindex_for_cellid dict. + + Tries to minimize the additional construction of the fileindex hashtable by spatial indexing. + ''' + if hasattr(self,"skipread"): + return + if coords.shape[0] == 0: + return + + if self.get_cellid_spatial_index() == None: + + self.__read_fileindex_for_cellid() + return + + # If the query bounding box volume is small, we really should use spatial indexing instead of loading the whole + # domain + volume_threshold = 0.01 + rank_ratio_threshold = 1./8. + + mins = np.min(coords, axis=0)-self.get_grid_epsilon() + maxs = np.max(coords, axis=0)+self.get_grid_epsilon() + query_window = np.array([mins[0],maxs[0],mins[1],maxs[1],mins[2],maxs[2]]) + + + if(coords.shape[0] == 1): + rankids = set(self.__cellid_spatial_index.intersection(query_window)) + rankids = rankids - self.__loaded_fileindex_ranks + for rankid in rankids: + self.get_cellid_locations_rank(rankid) + + + query_volume = np.prod(maxs-mins) + full_domain_extents = self.get_fsgrid_mesh_extent() + full_domain_mins = full_domain_extents[0:3] + full_domain_maxs = full_domain_extents[3:6] + full_domain_volume = np.prod(full_domain_maxs-full_domain_mins) + + if query_volume < volume_threshold*full_domain_volume: # If we are querying a small volume, it is fast to do a partial update. + rankids = set(self.__cellid_spatial_index.intersection(query_window)) + rankids = rankids - self.__loaded_fileindex_ranks + if(len(rankids)>0): + # print("single box") + for rankid in rankids: + self.get_cellid_locations_rank(rankid) + else: + if coords.shape[0] > 0.001*self.read_attribute(name="CellID",mesh="SpatialGrid",attribute="arraysize",tag="VARIABLE"): + # Guess that there are too many cellids to handle even the intersection search sensibly + self.__read_fileindex_for_cellid() + else: + query_windows = np.zeros((coords.shape[0],6)) + query_windows[:,0::2] = coords - self.get_grid_epsilon() + query_windows[:,1::2] = coords + self.get_grid_epsilon() + rankids = set() + for qw in query_windows: # There is a bulk intersection_v function but it does not work? + rankids.update(self.__cellid_spatial_index.intersection(qw)) + rankids = rankids - self.__loaded_fileindex_ranks + # print(len(rankids), "ranks to load") + if len(rankids) < self.get_numWritingRanks("SpatialGrid")*rank_ratio_threshold: + # print("few ranks") + for rankid in rankids: + self.get_cellid_locations_rank(rankid) + else: # Fallback: read everything + # print("Fallback") + self.__read_fileindex_for_cellid() + + def get_cellid(self, coords): ''' Returns the cell ids at given coordinates @@ -2407,7 +2815,12 @@ def get_cellid(self, coords): # If needed, read the file index for cellid # if len(self.__fileindex_for_cellid) == 0: - self.__read_fileindex_for_cellid() + if hasattr(self,"skipread"): + pass + else: + self.do_partial_fileindex_update(coordinates) + + #good_ids = self.read_variable("CellID") # good_ids = np.array(list(self.__fileindex_for_cellid.keys())) # good_ids.sort() @@ -2725,12 +3138,14 @@ def get_cell_corner_vertices(self, cids): ''' - mask = ~dict_keys_exist(self.__cell_vertices,cids,prune_unique=False) - coords = self.get_cell_coordinates(cids[mask]) - vertices = np.zeros((len(cids[mask]), 8, 3),dtype=int) + mask = ~dict_keys_exist(self.__cell_corner_vertices,cids,prune_unique=False) + cell_vertex_sets = {} if(len(cids[mask]) > 0): + coords = self.get_cell_coordinates(cids[mask]) + vertices = np.zeros((len(cids[mask]), 8, 3),dtype=int) + ii = 0 for x in [-1,1]: for y in [-1,1]: @@ -2747,7 +3162,7 @@ def get_cell_corner_vertices(self, cids): self.__cell_corner_vertices.update(cell_vertex_sets) - for i, c in enumerate(cids[~mask]): + for c in cids[~mask]: cell_vertex_sets[c] = self.__cell_corner_vertices[c] return cell_vertex_sets @@ -2757,20 +3172,23 @@ def get_cell_corner_vertices(self, cids): def build_cell_neighborhoods(self, cids): mask = ~dict_keys_exist(self.__cell_neighbours, cids, prune_unique=False) - cell_vertex_sets = self.get_cell_corner_vertices(cids[mask]) # these are enough to fetch the neighbours - - cell_neighbor_sets = {c: set() for c in cell_vertex_sets.keys()} - vertices_todo = set().union(*cell_vertex_sets.values()) - neighbor_tuples_dict = self.build_dual_from_vertices(list(vertices_todo)) - for c,verts in cell_vertex_sets.items(): - # neighbor_tuples = self.build_dual_from_vertices(verts) - # cell_neighbor_sets[c].update(set().union(*neighbor_tuples.values())) - cell_neighbor_sets[c].update(set().union(*itemgetter(*cell_vertex_sets[c])(neighbor_tuples_dict))) - self.__cell_neighbours.update(cell_neighbor_sets) + cell_neighbor_sets = {} + + if len(cids[mask]) > 0: + cell_vertex_sets = self.get_cell_corner_vertices(cids[mask]) # these are enough to fetch the neighbours + cell_neighbor_sets = {c: set() for c in cell_vertex_sets.keys()} + vertices_todo = set().union(*cell_vertex_sets.values()) + neighbor_tuples_dict = self.build_dual_from_vertices(list(vertices_todo)) + for c in cell_vertex_sets.keys(): + # neighbor_tuples = self.build_dual_from_vertices(verts) + # cell_neighbor_sets[c].update(set().union(*neighbor_tuples.values())) + cell_neighbor_sets[c].update(set().union(*itemgetter(*cell_vertex_sets[c])(neighbor_tuples_dict))) + + self.__cell_neighbours.update(cell_neighbor_sets) for c in cids[~mask]: - cell_neighbor_sets[c] = self.__cell_neighbors[c] + cell_neighbor_sets[c] = self.__cell_neighbours[c] return cell_neighbor_sets @@ -3335,16 +3753,20 @@ def read_parameter(self, name): .. seealso:: :func:`read_variable` :func:`read_variable_info` ''' - - # Special handling for time - if name=="time": - if self.check_parameter(name="t"): - return self.read(name="t", tag="PARAMETER") - if name=="t": - if self.check_parameter(name="time"): - return self.read(name="time", tag="PARAMETER") - return self.read(name=name, tag="PARAMETER") + if name in self.__params_cache.keys(): + return self.__params_cache[name] + else: + # Special handling for time + if name=="time": + if self.check_parameter(name="t"): + return self.read(name="t", tag="PARAMETER") + if name=="t": + if self.check_parameter(name="time"): + return self.read(name="time", tag="PARAMETER") + val = self.read(name=name, tag="PARAMETER") + self.__params_cache[name] = val + return val def read_velocity_distribution_dense(self, cellid, pop="proton", regularize=True, setThreshold=None): ''' @@ -3866,4 +4288,116 @@ def optimize_clear_fileindex_for_cellid(self): ''' self.__fileindex_for_cellid = {} + def get_mesh_domain_extents(self, mesh): + if (mesh == 'fsgrid'): + raise NotImplementedError + elif (mesh == 'ionosphere'): + raise NotImplementedError + elif (mesh == 'SpatialGrid'): + MESH_DOMAIN_SIZES = self.read(name="", tag="MESH_DOMAIN_SIZES", mesh="SpatialGrid") + n_domain_cells = MESH_DOMAIN_SIZES[:,0]-MESH_DOMAIN_SIZES[:,1] + cids_all = self.read_variable("CellID") + lowcorners_all = self.read_variable("vg_coordinates_cell_lowcorner") + dxs_all = self.read_variable("vg_dx") + + + extents = np.zeros((MESH_DOMAIN_SIZES.shape[0],6)) + start = 0 + end = 0 + for i, ncells in enumerate(n_domain_cells): + end = start+ncells + cids = cids_all[start:end] + lowcorners = lowcorners_all[start:end,:] + dxs = dxs_all[start:end,:] + highcorners = lowcorners+dxs + mins = np.min(lowcorners,axis=0) + maxs = np.max(highcorners,axis=0) + + extents[i,:] = [mins[0],maxs[0],mins[1],maxs[1],mins[2],maxs[2]] + + start = end + + return extents.reshape((-1),order="C") + + else: + raise ValueError + + def get_cache_folder(self): + + fn = self.file_name + head,tail = os.path.split(fn) + path = head + numslist = re.findall(r'\d+(?=\.vlsv)', tail) + + if(len(numslist) == 0): + path = os.path.join(path,"vlsvcache",tail[:-5]) + else: + nums = numslist[-1] + head, tail = tail.split(nums) + + leading_zero = True + path = os.path.join(path,"vlsvcache",head[:-1]) + for i,n in enumerate(nums): + if n == '0' and leading_zero: continue + if leading_zero: + fmt = "{:07d}" + else: + fmt = "{:0"+str(7-i)+"d}" + path = os.path.join(path, fmt.format(int(n)*10**(len(nums)-i-1))) + leading_zero = False + + # print(path) + return path + + def cache_neighbor_stencils(self): + self.load_neighbor_stencils_from_filecache() + path = self.get_cache_folder() + os.makedirs(path,exist_ok=True) + cache_file_neighbors = os.path.join(path, "neighbors_cache.pkl") + with open(cache_file_neighbors,'wb') as cache: + pickle.dump({ + "cell_neighbours": self.__cell_neighbours, + "cell_vertices":self.__cell_vertices, + "cell_corner_vertices":self.__cell_corner_vertices} + ,cache) + + def load_neighbor_stencils_from_filecache(self): + if self.__neighbors_cache_available and not self.__neighbors_cache_loaded: + path = self.get_cache_folder() + cache_file_neighbors = os.path.join(path, "neighbors_cache.pkl") + if(os.path.isfile(cache_file_neighbors)): + with open(cache_file_neighbors,'rb') as cache: + loaded = pickle.load(cache) + + self.__cell_neighbours.update(loaded["cell_neighbours"]) + self.__cell_vertices.update(loaded["cell_vertices"]) + self.__cell_corner_vertices.update(loaded["cell_corner_vertices"]) + self.__neighbors_cache_len = len(self.__cell_neighbours) + else: + self.__neighbors_cache_loaded = True + + def set_cellid_spatial_index(self, force=False): + self.__cellid_spatial_index = self.__metadata_cache.set_cellid_spatial_index(force) + + def get_cellid_spatial_index(self, force=False): + if not force: + if self.__cellid_spatial_index is None: + self.__cellid_spatial_index = self.__metadata_cache.get_cellid_spatial_index(force) + else: + pass + else: + self.__cellid_spatial_index = self.__metadata_cache.set_cellid_spatial_index(force) + + return self.__cellid_spatial_index + + def clear_cache_folder(self): + path = self.get_cache_folder() + import shutil + shutil.rmtree(path) + + def cache_optimization_files(self, force=False): + ''' Create cached optimization files for this reader object (e.g. spatial index) + + ''' + self.__metadata_cache.set_cellid_spatial_index(force) \ No newline at end of file diff --git a/analysator/pyVlsv/vlsvvtkinterface.py b/analysator/pyVlsv/vlsvvtkinterface.py index f54a0612..8663e692 100644 --- a/analysator/pyVlsv/vlsvvtkinterface.py +++ b/analysator/pyVlsv/vlsvvtkinterface.py @@ -24,8 +24,10 @@ import logging import numpy as np import cProfile -import os.path +import os import pickle +from operator import itemgetter +import numbers import analysator as pt try: @@ -627,6 +629,7 @@ def __init__(self): self.__cellarrays = [] self._arrayselection = vtk.vtkDataArraySelection() self._arrayselection.AddObserver("ModifiedEvent", createModifiedCallback(self)) + self.__cellIDtoIdx = {} def SetFileName(self, filename): if filename != self.__FileName: @@ -635,7 +638,7 @@ def SetFileName(self, filename): if self.__FileName is not None: self.__reader = pt.vlsvfile.VlsvReader(self.__FileName) fn = os.path.basename(self.__FileName) - self.__metafile = os.path.join(os.path.dirname(self.__FileName),"vlsvmeta",fn[:-5]+"_meta") + self.__metafile = os.path.join(self.__reader.get_cache_folder(),"vlsvvtkcache.pkl") def GetFileName(self): return self.__FileName @@ -716,6 +719,8 @@ def children(cid, level): def getDescriptor(self, reinit=False): + if hasattr(self, "__descriptor") and hasattr(self, "__idxToFileIndex"): + return self.__descriptor, self.__idxToFileIndex try: if reinit: raise Exception("reinit = True") @@ -728,13 +733,24 @@ def getDescriptor(self, reinit=False): print("Re-initializing HTG, no metadata accessible because of ", e) self.__descriptor, self.__idxToFileIndex = self.buildDescriptor() if not os.path.isdir(os.path.dirname(self.__metafile)): - os.mkdir(os.path.dirname(self.__metafile)) + os.makedirs(os.path.dirname(self.__metafile), exist_ok=True) with open(self.__metafile,'wb') as mfile: pickle.dump({"descr":self.__descriptor, "idxToFileIndexMap":self.__idxToFileIndex}, mfile) return self.__descriptor, self.__idxToFileIndex + + def getCellIDtoIdxMap(self): + if len(self.__cellIDtoIdx) == 0: + FileIndexToID = {v:k for k,v in self.getDescriptor()[1].items()} + for c,fi in self.__reader._VlsvReader__fileindex_for_cellid.items(): + self.__cellIDtoIdx[c] = FileIndexToID[fi] + + return self.__cellIDtoIdx def getHTG(self): + if self.__htg is not None: + return self.__htg + f = self.__reader descr, idxToFileIndex = self.getDescriptor() @@ -838,7 +854,7 @@ def RequestUpdateExtent(self, request, inInfo, outInfo): ''' def addArrayFromVlsv(self, varname): - htg = self.__htg + htg = self.getHTG() # Do not re-add an already existing array if htg.GetCellData().HasArray(varname): print("skipped existing array") @@ -919,9 +935,16 @@ def ReadAllScalarsOn(self): # self.__htg.addArrayFromVlsv(name) - - def GetDataArraySelection(self): - return self._arrayselection +# TODO + # def GetDataArraySelection(self): + # return self._arrayselection + + # def GetIdx(self, cellid): + # cidToIdx = self.getCellIDtoIdxMap() + # if isinstance(cellid,numbers.Number): + # return cidToIdx[cellid] + # else: + # return itemgetter(*cellid)(cidToIdx) def __main__(): diff --git a/analysator/pyVlsv/vlsvwriter.py b/analysator/pyVlsv/vlsvwriter.py index 9d206224..330e3e51 100644 --- a/analysator/pyVlsv/vlsvwriter.py +++ b/analysator/pyVlsv/vlsvwriter.py @@ -29,7 +29,6 @@ import os import warnings from reduction import datareducers,data_operators -import warnings class VlsvWriter(object): ''' Class for reading VLSV files @@ -112,17 +111,24 @@ def __initialize( self, vlsvReader, copy_meshes=None ): tags['MESH_OFFSETS'] = '' tags['MESH'] = '' tags['MESH_DOMAIN_SIZES'] = '' + tags['MESH_DOMAIN_EXTENTS'] = '' tags['MESH_GHOST_DOMAINS'] = '' tags['MESH_GHOST_LOCALIDS'] = '' tags['CellID'] = '' tags['MESH_BBOX'] = '' tags['COORDS'] = '' + xml_mesh_tags = {} # Copy the xml root for child in xml_root: if child.tag in tags: if 'name' in child.attrib: name = child.attrib['name'] else: name = '' + if 'mesh' in child.attrib: + mesh = child.attrib['mesh'] + xml_mesh_tags.setdefault(mesh, []).append(child.tag) + else: + mesh = None if 'mesh' in child.attrib: mesh = child.attrib['mesh'] else: mesh = '' tag = child.tag @@ -142,6 +148,15 @@ def __initialize( self, vlsvReader, copy_meshes=None ): # Write the data: self.__write( data=data, name=name, tag=tag, mesh=mesh, extra_attribs=extra_attribs ) + + # Find out and write possibly nonexisting metadata + for mesh, tags in xml_mesh_tags.items(): + if mesh == "SpatialGrid": + if "MESH_DOMAIN_EXTENTS" not in tags: + extents = vlsvReader.get_mesh_domain_extents(mesh) + # print(extents) + self.__write( data = extents, name='', tag="MESH_DOMAIN_EXTENTS", mesh=mesh) + def copy_variables( self, vlsvReader, varlist=None ): diff --git a/examples/setup-bulk-caches.py b/examples/setup-bulk-caches.py new file mode 100644 index 00000000..6884f5e8 --- /dev/null +++ b/examples/setup-bulk-caches.py @@ -0,0 +1,17 @@ +import analysator as pt +import sys +import glob + +if len(sys.argv) != 2: + print("Need globbable argument, eg. python setup-bulk-caches.py /wrk-vakka/group/spacephysics/vlasiator/3D/FHA/bulk1/bulk1.*.vlsv") + sys.exit(1) + +fns = glob.glob(sys.argv[1]) + +fns.sort() + + +for fn in fns: + f = pt.vlsvfile.VlsvReader(fn) + f.cache_optimization_files(True) + del f diff --git a/pyproject.toml b/pyproject.toml index 5da7dfc5..705e5e83 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,22 +26,20 @@ dependencies = [ "matplotlib", "packaging", "scikit-image", + "vtk>=9.2", + "yt", + "rtree" ] [project.optional-dependencies] none = [ ] -vtk = [ - "vtk>=9.2", -] all = [ - "analysator[vtk]", ] bvtk = [ "vtk==9.2.6", ] testpackage = [ "opencv-python", - "analysator[vtk]", ] [project.urls] Homepage = "https://github.com/fmihpc/analysator" diff --git a/requirements.txt b/requirements.txt index 444e84df..7f8ffdcc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,3 +4,4 @@ matplotlib scikit-image packaging vtk>=9.2 +rtree diff --git a/scripts/magnetopause.py b/scripts/magnetopause.py new file mode 100644 index 00000000..94ce29d2 --- /dev/null +++ b/scripts/magnetopause.py @@ -0,0 +1,233 @@ +"""Functions for finding the magnetopause from .vlsv data + + Example use with own arguments passed to streamline function: + + .. code-block:: python + + datafile = "bulkfile.vlsv" + vtpoutfilen = "magnetopause.vtp" + vlsvoutfilen = "magnetopause.vlsv" + + SW_args = {"seeds_n":300, "seeds_x0":150e6, "seeds_range":[-5*6371000, 5*6371000], + "dl":5e5, "iterations":int(1500), "end_x":-50*6371000, "x_point_n":100, "sector_n":36*2} + + surface, SDF = magnetopause(datafile, + method="streamlines", + return_SDF=True, + return_surface=True, + method_args=SW_args) + + write_vtk_surface_to_file(surface, vtpoutfilen) + write_SDF_to_file(SDF, datafile, vlsvoutfilen) + + + Example use for 0.0-0.4 beta* magnetopause with tetrahedrons with sides longer than 1Re excluded (aplha = 1Re): + + .. code-block:: python + + datafile = "bulkfile.vlsv" + vtpoutfilen = "magnetopause.vtp" + + surface, __ = magnetopause(datafile, + method="beta_star_with_connectivity", + beta_star_range=[0.0, 0.4], + Delaunay_alpha=1*6371000, + return_SDF=False, + return_surface=True) + + write_vtk_surface_to_file(surface, vtpoutfilen) + + +""" + +import analysator as pt +import numpy as np +import vtk +from scripts import shue, regions +import logging + + +R_E = 6371000 + + +def write_vtk_surface_to_file(vtkSurface, outfilen): + writer = vtk.vtkXMLPolyDataWriter() + writer.SetInputConnection(vtkSurface.GetOutputPort()) + writer.SetFileName(outfilen) + writer.Write() + logging.info("wrote ", outfilen) + +def write_SDF_to_file(SDF, datafilen, outfilen): + f = pt.vlsvfile.VlsvReader(file_name=datafilen) + writer = pt.vlsvfile.VlsvWriter(f, outfilen) + writer.copy_variables_list(f, ["CellID"]) + writer.write_variable_info(pt.calculations.VariableInfo(SDF, "SDF_magnetopause", "-", latex="",latexunits=""),"SpatialGrid",1) + logging.info("wrote ", outfilen) + + + + +def magnetopause(datafilen, method="beta_star_with_connectivity", own_tresholds=None, return_surface=True, return_SDF=True, SDF_points=None, Delaunay_alpha=None, beta_star_range=[0.4, 0.5], method_args={}): # TODO: separate streamline suface and vtkDelaunay3d surface in streamline method + """Finds the magnetopause using the specified method. Surface is constructed using vtk's Delaunay3d triangulation which results in a convex hull if no Delaunay_alpha is given. + Returns vtk.vtkDataSetSurfaceFilter object and/or signed distances (negative -> inside magnetopause) (=SDF) to all cells + Note that using alpha for Delaunay might make SDF different from expected inside the magnetosphere, especially if surface is constructed with points not everywhere in the magnetosphere (e.g. beta* 0.4-0.5) or if simulation grid size is larger than alpha + + :param datafilen: a .vlsv bulk file name (and path) + :kword method: str, default "beta_star_with_connectivity", other options "beta_star", "streamlines", "shue", "dict" + :kword own_tresholds: if using method "dict", a dictionary with conditions for the magnetopause/magnetosphere must be given where key is data name in datafile and value is treshold used (see treshold_mask()) + :kword return_surface: True/False, return vtkDataSetSurfaceFilter object + :kword return_SDF: True/False, return array of distances in m to SDF_points in point input order, negative distance inside the surface + :kword SDF_points: optionally give array of own points to calculate signed distances to. If not given, distances will be to cell centres in the order of f.read_variable("CellID") output + :kword Delaunay_alpha: alpha (float) to give to vtkDelaunay3d, None -> convex hull, alpha=__: surface egdes longer than __ will be excluded + :kword beta_star_range: [min, max] treshold rage to use with methods "beta_star" and "beta_star_with_connectivity" + :kword method_args: dict of keyword arguments to be passed down to external functions (for streamlines and shue) + :returns: vtkDataSetSurfaceFilter object of convex hull or alpha shape if return_surface=True, signed distance field of convex hull or alpha shape of magnetopause if return_SDF=True + """ + + f = pt.vlsvfile.VlsvReader(file_name=datafilen) + cellids = f.read_variable("CellID") + [xmin, ymin, zmin, xmax, ymax, zmax] = f.get_spatial_mesh_extent() + query_points = f.get_cell_coordinates(cellids) # for SDF, all centre points of cells + + if method == "streamlines": + + if "streamline_seeds" in method_args or "seeds_n" in method_args: + vertices, manual_vtkSurface = pt.calculations.find_magnetopause_sw_streamline_3d(datafilen, **method_args) + else: + seeds_x0=150e6 + dl=5e5 + iters = int(((seeds_x0-xmin)/dl)+100) + sector_n = 36*3 + vertices, manual_vtkSurface = pt.calculations.find_magnetopause_sw_streamline_3d(datafilen, seeds_n=300, seeds_x0=seeds_x0, seeds_range=[-5*6371000, 5*6371000], + dl=dl, iterations=iters, end_x=xmin+10*6371000, x_point_n=200, sector_n=sector_n) + + # make the magnetopause surface from vertice points + np.random.shuffle(vertices) # helps Delaunay triangulation + vtkSurface, SDF = regions.vtkDelaunay3d_SDF(query_points, vertices, Delaunay_alpha) + + elif method == "beta_star": + # magnetopause from beta_star only + mpause_flags = regions.treshold_mask(f.read_variable("vg_beta_star"), beta_star_range) + contour_coords = f.get_cell_coordinates(cellids[mpause_flags!=0]) + + # make a convex hull surface with vtk's Delaunay + np.random.shuffle(contour_coords) + vtkSurface, SDF = regions.vtkDelaunay3d_SDF(query_points, contour_coords, Delaunay_alpha) + + + + elif method == "beta_star_with_connectivity": + # magnetopause from beta_star, with connectivity if possible + betastar_region = regions.treshold_mask(f.read_variable("vg_beta_star"), beta_star_range) + try: + connectivity_region = regions.treshold_mask(f.read_variable("vg_connection"), 0) # closed-closed magnetic field lines + magnetosphere_proper = np.where((connectivity_region==1) | (betastar_region==1), 1, 0) + contour_coords = f.get_cell_coordinates(cellids[magnetosphere_proper==1]) + np.save("pointcloud.npy", contour_coords) + except: + logging.warning("using field line connectivity for magnetosphere did not work, using only beta*") + #condition_dict = {"beta_star": [0.5, 0.6]} # FIC: [0.4, 0.5]) # EGE: [0.9, 1.0]) # max 0.6 in FHA to not take flyaways from outside magnetopause + mpause_flags = np.where(betastar_region==1, 1, 0) + contour_coords = f.get_cell_coordinates(cellids[mpause_flags!=0]) + + # make a convex hull surface with vtk's Delaunay + vtkSurface, SDF = regions.vtkDelaunay3d_SDF(query_points, contour_coords, Delaunay_alpha) + + #elif method == "beta_star_with_fieldlines": # either incredibly slow or does not work, don't use without fixing #TODO proprer measure of actual field line backwall point + # # magnetopause from beta_star, with field lines connecting to back wall if possible + # betastar_region = regions.treshold_mask(f.read_variable("vg_beta_star"), beta_star_range) + # try: + # fw_region = regions.treshold_mask(f.read_variable("vg_connection_coordinates_fw")[:,0], [None, xmin+5e7]) + # bw_region = regions.treshold_mask(f.read_variable("vg_connection_coordinates_bw")[:,0], [None, xmin+5e7]) + # #connectivity_region = regions.treshold_mask(f.read_variable("vg_connection"), 0) # closed-closed magnetic field lines + # magnetosphere_proper = np.where(((fw_region==1) | (betastar_region==1) | (bw_region==1)), 1, 0) + # contour_coords = f.get_cell_coordinates(cellids[magnetosphere_proper==1]) + # except: + # logging.warning("using field lines for magnetosphere did not work, using only beta*") + # #condition_dict = {"beta_star": [0.5, 0.6]} # FIC: [0.4, 0.5]) # EGE: [0.9, 1.0]) # max 0.6 in FHA to not take flyaways from outside magnetopause + # mpause_flags = np.where(betastar_region==1, 1, 0) + # contour_coords = f.get_cell_coordinates(cellids[mpause_flags!=0]) + # + # # make a convex hull surface with vtk's Delaunay + # vtkSurface, SDF = regions.vtkDelaunay3d_SDF(query_points, contour_coords, Delaunay_alpha) + + + + elif method == "dict": + variable_dict = {} + for key in own_tresholds: + variable_dict[key] = f.read_variable(key, cellids=-1) + + # same method as beta* but from user-defined variables as initial contour + flags = regions.make_region_flags(variable_dict, own_tresholds, flag_type="01") + treshold_coords = f.get_cell_coordinates(cellids[flags!=0]) + vtkSurface, SDF = regions.vtkDelaunay3d_SDF(query_points, treshold_coords, Delaunay_alpha) + + + elif method == "shue": + # note: might not be correct but should produce something, should recheck the projection coordinates + + theta = np.linspace(0, 2*np.pi/3 , 200) # magnetotail length decided here by trial and error: [0, 5*np.pi/6] ~ -350e6 m, [0, 2*np.pi/3] ~ -100e6 m in EGE + if "run" in method_args or "B_z" in method_args: + r, __, __ = shue.f_shue(theta, **method_args) + else: + r, __, __ = shue.f_shue(theta, B_z = -10, n_p = 1, v_sw = 750) # for runs not in shue.py B_z, n_p, and v_sw need to be specified and run=None + + # 2d one-sided magnetopause + xs = r*np.cos(theta) + ys = r*np.sin(theta) + + # 3d projection for complete magnetopause + psis = np.linspace(0, 2*np.pi, 100) + coords = np.zeros((len(theta)*len(psis), 3)) + i=0 + for x,y in zip(xs,ys): + for psi in psis: + coords[i] = np.array([x, y*np.sin(psi), y*np.cos(psi)]) + i += 1 + + coords = coords*R_E + + # surface and SDF + np.random.shuffle(coords) # helps Delaunay triangulation + vtkSurface, SDF = regions.vtkDelaunay3d_SDF(query_points, coords) # alpha does nothing here + + + + else: + logging.error("Magnetopause method not recognized. Use one of the options: \"beta_star\", \"beta_star_with_connectivity\", \"streamlines\", \"shue\", \"dict\"") + exit() + + if return_surface and return_SDF: + return vtkSurface, SDF + elif return_surface: + return vtkSurface, None + elif return_SDF: + return None, SDF + + # beta* + B connectivity: + #connectivity_region = regions.treshold_mask(f.read_variable("vg_connection"), 0) + #betastar_region = regions.treshold_mask(f.read_variable("beta_star"), [0.0, 0.5]) + #magnetosphere_proper = np.where((connectivity_region==1) | (betastar_region==1), 1, 0) + + +def main(): + + datafile = "/wrk-vakka/group/spacephysics/vlasiator/3D/EGE/bulk/bulk.0002000.vlsv" + vtpoutfilen = "EGE_magnetopause_t2000.vtp" + vlsvoutfilen = "EGE_magnetopause_t2000.vlsv" + + surface, SDF = magnetopause(datafile, + method="beta_star", + beta_star_range=[0.9, 1.0], + return_SDF=True, + return_surface=True) + + write_vtk_surface_to_file(surface, vtpoutfilen) + write_SDF_to_file(SDF, datafile, vlsvoutfilen) + + +if __name__ == "__main__": + + main() + diff --git a/scripts/magnetopause2d.py b/scripts/magnetopause2d.py deleted file mode 100644 index 5bf4e5df..00000000 --- a/scripts/magnetopause2d.py +++ /dev/null @@ -1,185 +0,0 @@ -''' -Finds the magnetopause position by tracing steamines of the plasma flow for two-dimensional Vlasiator runs. Needs the yt package. -''' - -import numpy as np -import analysator as pt -import plot_colormap -import yt -from yt.visualization.api import Streamlines - - - - -def interpolate(streamline, x_points): - - arr = np.array(streamline) - - # set arrays for interpolation - xp = arr[:,0][::-1] - zp = arr[:,2][::-1] - - #interpolate z coordinates - z_points = np.interp(x_points, xp, zp, left=np.NaN, right=np.NaN) - - return np.array([x_points, z_points]) - - -def make_streamlines(vlsvFileName): - ## make streamlines - boxre = [0,0] - - # bulk file - f = pt.vlsvfile.VlsvReader(file_name=vlsvFileName) - - #get box coordinates from data - [xmin, ymin, zmin, xmax, ymax, zmax] = f.get_spatial_mesh_extent() - [xsize, ysize, zsize] = f.get_spatial_mesh_size() - simext_m =[xmin,xmax,ymin,ymax,zmin,zmax] - sizes = np.array([xsize,ysize,zsize]) - - def to_Re(m): #meters to Re - return m/6371000 - - simext = list(map(to_Re, simext_m)) - boxcoords=list(simext) - - cellids = f.read_variable("CellID") - - #Read the data from vlsv-file - Vx = f.read_variable("v", operator="x") - Vz = f.read_variable("v", operator="z") - - - #Re-shape variable data - Vxs=Vx[np.argsort(cellids)].reshape(f.get_spatial_mesh_size(), order="F") - Vys = np.zeros_like(Vxs) - Vzs=Vz[np.argsort(cellids)].reshape(f.get_spatial_mesh_size(), order="F") - - data=dict(Vx=Vxs,Vy=Vys,Vz=Vzs) - - #Create streamline seeds (starting points for streamlines) - seedN = 200 # number of streamlines wanted - streamline_seeds = np.array([[20, 0 ,i] for i in np.linspace(-4, 4, seedN)]) - - #streamline_seeds = np.array(streamline_seeds) - #dataset in yt-form - yt_dataset = yt.load_uniform_grid( - data, - sizes, - bbox=np.array([[boxcoords[0], boxcoords[1]], - [boxcoords[2],boxcoords[3]], - [boxcoords[4],boxcoords[5]]])) - - - #data, seeds, dictionary positions, lenght of lines - streamlines_pos = yt.visualization.api.Streamlines(yt_dataset, streamline_seeds, - "Vx", "Vy", "Vz", length=40, direction=1) - - #where the magic happens - streamlines_pos.integrate_through_volume() - return np.array(streamlines_pos.streamlines) - -def make_magnetopause(streams): - - streampoints = np.reshape(streams, (streams.shape[0]*streams.shape[1], 3)) #all the points in one array - - ## find the subsolar dayside point in the positive x-axis - ## do this by finding a stremline point on positive x axis closest to the Earth - x_axis_points = streampoints[np.floor(streampoints[:,2])==0] - x_axis_points[x_axis_points<0] = 800 - subsolar_x =np.min(x_axis_points[:,0]) - - ## define points in the x axis where to find magnetopause points on the yz-plane - x_points = np.arange(subsolar_x, -10, -0.2) - - ## interpolate more exact points for streamlines at exery x_point - new_streampoints = np.zeros((len(x_points), len(streams), 1)) # new array for keeping interpolated streamlines in form streamlines_new[x_point, streamline, z-coordinate] - i=0 - for stream in streams: - interpolated_streamline = interpolate(stream, x_points) - for j in range(0, len(x_points)): - new_streampoints[j, i,:] = interpolated_streamline[1,j] - i += 1 - - - ## start making the magnetopause - ## in each x_point, find the closest streamline to x-axis in the positive and negative z-axis - - pos_z_mpause = np.zeros((len(x_points), 2)) - neg_z_mpause = np.zeros((len(x_points), 2)) - - for i, x_point in enumerate(x_points): - pos = new_streampoints[i, new_streampoints[i,:] > 0] - neg = new_streampoints[i, new_streampoints[i,:] < 0] - - if (pos.size == 0) or (neg.size == 0): - raise ValueError('No streamlines found for x axis point, try adding streamlines or checking the x_points') - - # find points closest to x-axis and save found points - pos_z_mpause[i] = [x_point, pos[pos.argmin()]] - neg_z_mpause[i] = [x_point, neg[neg.argmax()]] - - magnetopause = np.concatenate((pos_z_mpause[::-1], np.array([[subsolar_x, 0]]), neg_z_mpause)) - - return magnetopause - - -def polar_to_cartesian(r, phi): - phi = np.deg2rad(phi) - y = r * np.cos(phi) - z = r * np.sin(phi) - return(y, z) - - - -def main(): - - ## get bulk data - run = 'BFD' - num = '2000' - - fileLocation="/wrk-vakka/group/spacephysics/vlasiator/2D/"+run+"/bulk/" - fileN = "bulk.000"+num+".vlsv" - - - ## STREAMLINES - streams = make_streamlines(fileLocation+fileN) - - ## MAGNETOPAUSE - magnetopause = make_magnetopause(streams) - - - ## PLOTTING - outdir="" - - # plot the magnetopause (and streamlines if needed) on top of colormap - def external_plot(ax,XmeshXY=None, YmeshXY=None, pass_maps=None): - plot_magnetopause = True - plot_streamlines = False - - if plot_streamlines: - for stream in streams: - ax.plot(stream[:,0], stream[:,2], color='paleturquoise', alpha=0.2) - - if plot_magnetopause: - ax.plot(magnetopause[:,0], magnetopause[:,1], color='cyan', linewidth=1.0) - - - plot_colormap.plot_colormap( - filename=fileLocation+fileN, - outputdir=outdir, - nooverwrite=None, - var="rho", #var="beta_star", - boxre = [-21,21,-21,21], - title=None, - draw=None, usesci=True, Earth=True, - external = external_plot, - run=run, - colormap='inferno', - ) - - - -if __name__ == "__main__": - main() diff --git a/scripts/magnetopause3d.py b/scripts/magnetopause3d.py deleted file mode 100644 index 906880d7..00000000 --- a/scripts/magnetopause3d.py +++ /dev/null @@ -1,376 +0,0 @@ -''' -Finds the magnetopause position by tracing steamines of the plasma flow for three-dimensional Vlasiator runs. Needs the yt package. -''' -from pyCalculations import ids3d -import matplotlib.pyplot as plt -import numpy as np -import analysator as pt -import plot_colormap3dslice -import yt -import math -from mpl_toolkits import mplot3d -from yt.visualization.api import Streamlines - - - - -def to_Re(m): #meters to Re - return m/6371000 - - -def cartesian_to_polar(cartesian_coords): # for segments of plane - y,z = cartesian_coords[0], cartesian_coords[1] - r = np.sqrt(z**2 + y**2) - phi = np.arctan2(z, y) - phi = np.rad2deg(phi) #angle in degrees - if phi < 0: phi = phi+360 - return(r, phi) - -def polar_to_cartesian(r, phi): - phi = np.deg2rad(phi) - y = r * np.cos(phi) - z = r * np.sin(phi) - return(y, z) - - -def interpolate(streamline, x_points): - arr = np.array(streamline) - - # set arrays for interpolation - xp = arr[:,0][::-1] - yp = arr[:,1][::-1] - zp = arr[:,2][::-1] - - #interpolate z from xz-plane - z_points = np.interp(x_points, xp, zp, left=np.NaN, right=np.NaN) - #interpolate y from xy-plane - y_points = np.interp(x_points, xp, yp, left=np.NaN, right=np.NaN) - - - return np.array([x_points, y_points, z_points]) - - - - -def make_surface(coords): - - ''' - Defines surface constructed of input coordinates as triangles - Returns list of verts and vert indices of surface triangles - - coordinates must be in form [...[c11, c21, c31, ... cn1],[c12, c22, c32, ... cn2],... - where cij = [xij, yij, zij], i marks slice, j marks yz-plane (x_coord) index - - How it works: - Three points make a triangle, triangles make the surface. - For every two planes next to each other: - - take every other point from plane1, every other from plane2 (in order!) - - from list of points: every three points closest to each other make a surface - - Example: - plane 1: [v1, v2, v3, v4] - plane 2: [v5, v6, v7, v8] - - -> list: [v1, v5, v2, v6, v3,...] - -> triangles: - v1 v5 v2 - v5 v2 v6 - v2 v6 v3 - . - . - . - - ''' - verts = [] #points - faces = [] #surface triangles - - slices_in_plane = len(coords[0]) - planes = len(coords) - - #get points - for plane in coords: - for vert in plane: - verts.append(vert) - - #get triangle (face) indices - #Let's consider the area between the first two planes - #ind1, ind2, ind3 for triangle indices - ind1 = 0 #starts from plane 1 - ind2 = slices_in_plane #starts from plane 2 - ind3 = 1 #starts from plane 1 - first_triangles = [] - - while len(first_triangles) < (2*slices_in_plane): - first_triangles.append([ind1, ind2, ind3]) - ind1 = ind2 - ind2 = ind3 - if (ind3 == (slices_in_plane*2)-1): #max index, go over to plane 1 first index - ind3 = 0 - elif (ind3 == 0): #last round, go to plane 2 first index - ind3 = slices_in_plane - else: - ind3 = ind1 + 1 - - first_triangles = np.array(first_triangles) - - #Now the rest of the triangles are just the first triangles + (index of area * slices_in_plane) - for area_index in range(planes-1): - next_triangles = [x + slices_in_plane*area_index for x in first_triangles] - faces.extend(next_triangles) - - return verts, faces - - - -def make_streamlines(vlsvFileName): - ## make streamlines - boxre = [0] - - # bulk file - f = pt.vlsvfile.VlsvReader(file_name=vlsvFileName) - - #get box coordinates from data - [xmin, ymin, zmin, xmax, ymax, zmax] = f.get_spatial_mesh_extent() - [xsize, ysize, zsize] = f.get_spatial_mesh_size() - [xsizefs, ysizefs, zsizefs] = f.get_fsgrid_mesh_size() - simext_m =[xmin,xmax,ymin,ymax,zmin,zmax] - sizes = np.array([xsize,ysize,zsize]) - sizesfs = np.array([xsizefs,ysizefs,zsizefs]) - - simext = list(map(to_Re, simext_m)) - boxcoords=list(simext) - - cellids = f.read_variable("CellID") - indexids = cellids.argsort() - cellids = cellids[indexids] - - reflevel = ids3d.refinement_level(xsize, ysize, zsize, cellids[-1]) - - - #Read the data from vlsv-file - V = f.read_variable("vg_v") - - #from m to Re - V = np.array([[to_Re(i) for i in j ]for j in V]) - - - V = V[indexids] - - if np.ndim(V)==1: - shape = None - elif np.ndim(V)==2: # vector variable - shape = V.shape[1] - elif np.ndim(V)==3: # tensor variable - shape = (V.shape[1], V.shape[2]) - - - Vdpoints = ids3d.idmesh3d2(cellids, V, reflevel, xsize, ysize, zsize, shape) - - - Vxs = Vdpoints[:,:,:,0] - Vys = Vdpoints[:,:,:,1] - Vzs = Vdpoints[:,:,:,2] - - data=dict(Vx=Vxs,Vy=Vys,Vz=Vzs) - - #Create streamline seeds (starting points for streamlines) - seedN = 50 #seeds per row, final seed count will be seedN*seedN ! - streamline_seeds = np.zeros([seedN**2, 3]) - - #range: np.arange(from, to, step) - t = np.linspace(-5, 5, seedN) - k = 0 - for i in t: - for j in t: - streamline_seeds[k] = [20, i, j] - k = k+1 - - - #dataset in yt-form - yt_dataset = yt.load_uniform_grid( - data, - sizesfs, - bbox=np.array([[boxcoords[0], boxcoords[1]], - [boxcoords[2],boxcoords[3]], - [boxcoords[4],boxcoords[5]]])) - - - # data, seeds, dictionary positions, length of streamlines - streamlines_pos = yt.visualization.api.Streamlines(yt_dataset, streamline_seeds, - "Vx", "Vy", "Vz", length=50, direction=1) - - # make the streamlines - streamlines_pos.integrate_through_volume() - - return np.array(streamlines_pos.streamlines) - - -def make_magnetopause(streams): - streampoints = np.reshape(streams, (streams.shape[0]*streams.shape[1], 3)) #all the points in one array - - ## find the subsolar dayside point in the x-axis - ## do this by finding a stremline point on positive x axis closest to the Earth - x_axis_points = streampoints[(np.floor(streampoints[:,1])==0) & (np.floor(streampoints[:,2])==0)] - subsolar_x =np.min(x_axis_points[:,0]) - - ## define points in the x axis where to find magnetopause points on the yz-plane - x_points = np.arange(subsolar_x, -15, -0.5) - - ## interpolate more exact points for streamlines at exery x_point - new_streampoints = np.zeros((len(x_points), len(streams), 2)) # new array for keeping interpolated streamlines in form new_streampoints[x_point, streamline, y and z -coordinates] - i=0 # streamline - for stream in streams: - interpolated_streamline = interpolate(stream, x_points) - - if type(interpolated_streamline) is np.ndarray: # don't use 'discarded' streamlines, see function interpolate() - for j in range(0, len(x_points)): - y,z = interpolated_streamline[1,j], interpolated_streamline[2,j] - new_streampoints[j, i,:] = np.array([y,z]) - i += 1 - - - - ## create a list of streamline points in polar coordinates (for every x_point) - polar_coords = np.zeros_like(new_streampoints) - for i in range(0,new_streampoints.shape[0]): - for j in range(0,new_streampoints.shape[1]): - polar_coords[i,j,:] = cartesian_to_polar(new_streampoints[i,j]) - - - ## now start making the magnetopause - ## in each x_point, divide the plane into sectors and look for the closest streamline to x-axis in the sector - sector_n = 36 - - ## if given sector number isn't divisible by 4, make it so because we want to have magnetopause points at exactly y=0 and z=0 for 2d slices of the whole thing - while sector_n%4 != 0: - sector_n +=1 - - sector_width = 360/sector_n - magnetopause = np.zeros((len(x_points), sector_n, 3)) - - for i,x_point in enumerate(x_points): #loop over every chosen x-axis point - # divide the yz-plane into sectors - for j, mean_sector_angle in enumerate(np.arange(0, 360, sector_width)): - min_angle = mean_sector_angle-sector_width/2 - max_angle = mean_sector_angle+sector_width/2 - - # find points that are in the sector - if mean_sector_angle == 0: # special case as the first sector needs streamlines around phi=0 - min_angle = min_angle+360 - # divide into phi<360 and phi>0 - sector1 = polar_coords[i, (polar_coords[i,:,1] <= 360)*(polar_coords[i,:,1] > min_angle)] - sector2 = polar_coords[i, (polar_coords[i,:,1] <= max_angle)*(polar_coords[i,:,1] >= 0)] - sector_points = np.concatenate((sector1, sector2)) - - else: - sector_points = polar_coords[i, (polar_coords[i,:,1] <= max_angle)*(polar_coords[i,:,1] > min_angle)] - - # discard 'points' with r=0 and check that there's at least one streamline point in the sector - sector_points = sector_points[sector_points[:,0] != 0.0] - if sector_points.size == 0: - raise ValueError('No streamlines found in the sector') - - # find the points closest to the x-axis - closest_point_radius = sector_points[sector_points[:,0].argmin(), 0] # smallest radius - - # return to cartesian coordinates and save as a magnetopause point at the middle of the sector - y,z = polar_to_cartesian(closest_point_radius, mean_sector_angle) - magnetopause[i,j,:] = [x_point, y, z] - - - # make a tip point for the magnetopause for prettier 3d plots - tip = np.array([subsolar_x, 0, 0]) - tips = np.tile(tip, (magnetopause.shape[1],1)) - magnetopause = np.vstack(([tips], magnetopause)) - - return magnetopause - - - - -def main(): - - ## get bulk data - run = 'EGI' - fileLocation="/wrk-vakka/group/spacephysics/vlasiator/3D/"+run+"/bulk/" - fileN = "bulk5.0000070.vlsv" - - ## STREAMLINES - streams = make_streamlines(fileLocation+fileN) - ## MAGNETOPAUSE - magnetopause = make_magnetopause(streams) - - - ## PLOTTING - outdir="" - - ## take separate arrays for different 2d slice plots - slices = magnetopause.shape[1] - quarter_slice = int(slices/4) - # xy plane: z=0 - xy_slice = np.concatenate((magnetopause[:,0][::-1], magnetopause[:,2*quarter_slice])) - # xz plane: y=0 - xz_slice = np.concatenate((magnetopause[:,quarter_slice][::-1], magnetopause[:,3*quarter_slice])) - - - #2D plots - # analysator 3dcolormapslice y=0 - if True: - def external_plot(ax,XmeshXY=None, YmeshXY=None, pass_maps=None, requestvariables=False): - if requestvariables==True: - return ['vg_v'] - ax.plot(xz_slice[:,0], xz_slice[:,2], color='limegreen', linewidth=1.5) - - - plot_colormap3dslice.plot_colormap3dslice( - filename=fileLocation+fileN, - outputdir=outdir, - run=run, - nooverwrite=None, - boxre = [-21,21,-21,21], - title=None, - draw=None, usesci=True, Earth=True, - external = external_plot, - colormap='inferno', - normal='y' - ) - - # analysator 3dcolormapslice z=0 - if True: - def external_plot(ax,XmeshXY=None, YmeshXY=None, pass_maps=None, requestvariables=False): - if requestvariables==True: - return ['vg_v'] - ax.plot(xy_slice[:,0], xy_slice[:,1], color='limegreen', linewidth=1.5) - - plot_colormap3dslice.plot_colormap3dslice( - filename=fileLocation+fileN, - outputdir=outdir, - run=run, - nooverwrite=None, - boxre = [-21,21,-21,21], - title=None, - draw=None, usesci=True, Earth=True, - external = external_plot, - colormap='inferno', - normal='z' - ) - - # 3D plot - # matplotlib 3d surface plot for single image - if False: - verts, faces = make_surface(magnetopause) - verts = np.array(verts) - - fig = plt.figure() - ax2 = fig.add_subplot(projection='3d') - ax2.plot_trisurf(verts[:, 0], verts[:,1], faces, verts[:, 2], linewidth=0.2, antialiased=True) - ax2.view_init(azim=-60, elev=5) - ax2.set_xlabel('X') - ax2.set_ylabel('Y') - ax2.set_zlabel('Z') - fig.tight_layout() - plt.savefig(outdir+run+'_3d_magnetopause.png') - - -if __name__ == "__main__": - main() diff --git a/scripts/plot_streamline_magnetopause_2d.py b/scripts/plot_streamline_magnetopause_2d.py new file mode 100644 index 00000000..4c17e8b0 --- /dev/null +++ b/scripts/plot_streamline_magnetopause_2d.py @@ -0,0 +1,41 @@ +""" +An example of using find_magnetopause_sw_streamline_2d() and plotting the output over colormap. +""" + +import analysator as pt + +def main(): + + run = 'BFD' # for output file name + file_name="/wrk-vakka/group/spacephysics/vlasiator/2D/BFD/bulk/bulk.0002000.vlsv" + + # magnetopause points + magnetopause = pt.calculations.find_magnetopause_sw_streamline_2d( + file_name, + streamline_seeds=None, + seeds_n=200, + seeds_x0=20*6371000, + seeds_range=[-5*6371000, 5*6371000], + streamline_length=45*6371000, + end_x=-15*6371000, + x_point_n=50) + + magnetopause = magnetopause*(1/6371000) # in RE + + def external_plot(ax,XmeshXY=None, YmeshXY=None, pass_maps=None): + ax.plot(magnetopause[:,0], magnetopause[:,1], color='cyan', linewidth=1.0) + + pt.plot.plot_colormap( + filename=file_name, + var="rho", + boxre = [-21,21,-21,21], + Earth=True, + external = external_plot, + run=run, + colormap='inferno', + ) + + +if __name__ == "__main__": + main() + diff --git a/scripts/regions.py b/scripts/regions.py new file mode 100644 index 00000000..7b7895f7 --- /dev/null +++ b/scripts/regions.py @@ -0,0 +1,506 @@ +"""Script and functions for creating sidecar files with SDF/region/boundary tags of plasma regions. + + Usage example where bow shock conditions are given to replace default conditions: + + .. code-block:: python + + datafile = "bulkfile.vlsv" + outfilen = "regions.vlsv" + RegionFlags(datafile, outfilen, regions=["all"], + region_conditions = {"bowshock": {"density": [2e6, None]}} + +""" + +import analysator as pt +import numpy as np +import vtk +from scripts import magnetopause +import logging + +R_E = 6371000 + +def vtkDelaunay3d_SDF(query_points, coordinates, alpha=None): + """Gives a signed distance to a convex hull or alpha shape surface created from given coordinates. + Note: if using alpha, SDF might not work, especially if the point cloud used for Delaunay does not fully cover e.g. lobes + + :param all_points: points ([x, y, z] coordinates in m) for which a signed distance to surface will be calculated + :param coordinates: coordinates (array of [x, y, z]:s in m) that are used to make a surface. + :kword alpha: alpha to be given to vtkDelaunay3D (e.g. R_E), removes surface edges that have length more than alpha so that the resulting surface is not convex. None -> convex hull + :returns: vtkDataSetSurfaceFilter() object, array of signed distances (negative sign: inside the surface, positive sign: outside the surface) + + """ + + points = vtk.vtkPoints()#.NewInstance() + for i in range(len(coordinates)): + points.InsertNextPoint(coordinates[i,0],coordinates[i,1],coordinates[i,2]) + polydata = vtk.vtkPolyData() + polydata.SetPoints(points) + polydata.Modified() + + # Delaunay convex hull + delaunay_3d = vtk.vtkDelaunay3D() + delaunay_3d.SetInputData(polydata) + if alpha is not None: + delaunay_3d.SetAlpha(alpha) + delaunay_3d.SetAlphaTets(1) + delaunay_3d.SetAlphaLines(0) + delaunay_3d.SetAlphaTris(0) + delaunay_3d.SetAlphaVerts(0) + delaunay_3d.Update() + + + data = delaunay_3d.GetOutput() + + surface = vtk.vtkDataSetSurfaceFilter() + surface.SetInputData(data) + surface.Update() + + implicitPolyDataDistance = vtk.vtkImplicitPolyDataDistance() + implicitPolyDataDistance.SetInput(surface.GetOutput()) + + convexhull_sdf = np.zeros(len(query_points)) + + # SDFs + for i,coord in enumerate(query_points): + convexhull_sdf[i] = implicitPolyDataDistance.EvaluateFunction(coord) + + return surface, convexhull_sdf + + + + +def treshold_mask(data_array, value): + """Chooses and flags cells where variable values match those given. If value is given as float, relative tolerance can be given + + :param data_array: array to mask, e.g. output of f.read_variable(name="proton/vg_rho", cellids=-1) + :param variable: str, variable name + :param value: value/values to use for masking; a float or int for exact match, a (value, relative tolerance) tuple, or [min value, max value] list pair where either can be None for less than or eq./more than or eq. value + :returns: 0/1 mask in same order as data_array, 1: variable value in array inside treshold values, 0: outside + """ + + if (data_array is None) or (value is None) or (np.isnan(data_array[0])): # either variable isn't usable or treshold has not been given + return None + + #mask = np.zeros((len(data_array))) + + if isinstance(value, float) or isinstance(value, int): # single value, exact + mask = np.where(np.isclose(data_array, value), 1, 0) + elif isinstance(value, tuple): # single value with tolerance + mask = np.where(np.isclose(data_array, value[0], rtol=value[1]), 1, 0) + else: + if value[0] is None: + mask = np.where((data_array <= value[1]), 1, 0) # anywhere where value is less than or equal to + elif value[1] is None: + mask = np.where((value[0] <= data_array), 1, 0) # anywhere where value is more than or equal to + else: + mask = np.where((value[0] <= data_array) & (data_array <= value[1]), 1, 0) # min/max + + if np.sum(mask[mask>0]) == 0: + logging.warning("Treshold mask didn't match any values in array") + return None + + return mask + + +def box_mask(f, coordpoints=None, cells=None, marginal=[150e6, 50e6, 50e6, 50e6, 50e6, 50e6]): + """Crops simulation box for calculations, output flags outside of cropped box will be 0 + + :param f: a VlsvReader + :kword coordpoints: coordinate points of cells to be masked + :kword cells: cellIDs to be masked + :kword marginal: 6-length list of wanted marginal lengths from mesh edges in meters [negx, negy, negz, posx, posy, posz] + :returns: Boolean mask in input order of coordinates/cells, -1 if no coorpoints or cells were given and mask could not be done + """ + [xmin, ymin, zmin, xmax, ymax, zmax] = f.get_spatial_mesh_extent() + + def is_inside(coords): + res = np.zeros((len(coords)), dtype=bool) + + for i,coord in enumerate(coords): + if np.any((coord[0] < xmin+marginal[0], coord[0] > xmax-marginal[3], + coord[1] < ymin+marginal[1], coord[1] > ymax-marginal[4], + coord[2] < zmin+marginal[2], coord[2] > zmax-marginal[5])): + res[i] = False + else: + res[i] =True + return res + + if coordpoints is not None: + return is_inside(coordpoints) + + elif cells is not None: + coords = f.get_cell_coordinates(cells) + return is_inside(coords) + + return -1 + +def write_flags(writer, flags, flag_name, mask=None): + """Writes flags into .vlsv-file with writer + + :param writer: a vlsvwriter + :param flags: an array of flags in order of cellids + :param flag_name: string, name of the flag variable + :kword mask: if flags are made from cropped cells, give mask used for cropping + """ + + if mask is not None: # flags exist only in cropped cells + new_flags = mask.astype(float) + new_flags[mask] = flags + writer.write_variable_info(pt.calculations.VariableInfo(new_flags, flag_name, "-", latex="",latexunits=""),"SpatialGrid",1) + + else: + #writer.write(flags,flag_name,'VARIABLE','SpatialGrid') + writer.write_variable_info(pt.calculations.VariableInfo(flags, flag_name, "-", latex="",latexunits=""),"SpatialGrid",1) + + logging.info(flag_name+" written to file") + + + +def make_region_flags(variable_dict, condition_dict, flag_type="01", mask=None): + """Makes region flags + + example: + + .. code-block:: python + + make_region_flags(variable_dict = {"density": f.read_variable(name="proton/vg_rho", cellids=-1)}, condition_dict = {"density": [None, 1e7]}) + + results in flag array in shape of read_variable output where cells where "proton/vg_rho" is less or equal to 1e7 are 1 and others are 0 + + + :param variable_dict: dictionary containing names and data arrays of variables, names must match those in condition_dict + :param condition_dict: dictionary containing names and conditions for variable data in variable_dict + :kword flag_type: default "01", optionally "fraction". "01" gives flags 1: all variables fill all conditions, 0: everything else; "fraction" flags are rounded fractions of varibles that fulfill conditions (e.g. flag 0.5 means one of two conditions was filled) + :kword mask: deafault None, optionally boolean array to use for variable dict arrays to search for a region from subset of cells (cells False in mask will be automatically flagged as 0) + + """ + + variable_flags = [] + + if mask is None: # search all cells + for key in condition_dict.keys(): + + try: # if key in variables + region = treshold_mask(variable_dict[key], condition_dict[key]) + if region is not None: variable_flags.append(region) + except: + logging.warning(key, "ignored") + + else: # search only masked area + for key in condition_dict.keys(): + try: # if key in variables + region = treshold_mask(variable_dict[key][mask], condition_dict[key]) + if region is not None: variable_flags.append(region) + except: + logging.warning(key, "ignored") + + + variable_flags = np.array(variable_flags) + + if len(variable_flags)==0: # no cells that fullfill conditions, hope that density is a key + if mask is None: + return np.zeros((len(variable_dict["density"]))) + else: + return np.zeros((len(variable_dict["density"][mask]))) + + + if flag_type == "01": # 1 only if all viable variables are 1, 0 otherwise + flags = np.where(np.sum(variable_flags, axis=0)==len(variable_flags), 1, 0) + return flags + + elif flag_type == "fraction": # rounded fraction of viable variables that are 1 + flags = np.sum(variable_flags, axis=0)/len(variable_flags) + return flags.round(decimals=2) + + + +def bowshock_SDF(f, variable_dict, query_points, own_condition_dict=None): + """Finds the bow shock by making a convex hull of either default variables/values based on upstream values or a user-defined variable/value dictionary. Returns signed distances from the convex hull to given query_points. + + :param f: a VlsvReader + :param variable_dict: dictionary containing names and variable arrays needed + :param query_points: xyz-coordinates of all points where the SDF will be calculated ([x1, y1, z1], [x2, y2, z2], ...) + :kword own_condition_dict: optional, dictionary with string variable names as keys and tresholds as values to pass to treshold_mask()-function + """ + + cellids =f.read_variable("CellID") + + if own_condition_dict is None: + # bow shock from upstream rho + # upstream point + upstream_point = [150e6, 0.0, 0.0] + upstream_cellid = f.get_cellid(upstream_point) + upstream_rho = f.read_variable(name="proton/vg_rho", cellids=upstream_cellid) + + bowshock_conditions = {"density": (1.5*upstream_rho, 0.1)} + bowshock_rho_flags = make_region_flags(variable_dict, bowshock_conditions, flag_type="01") + __, bowshock_rho_SDF = vtkDelaunay3d_SDF(query_points, f.get_cell_coordinates(cellids[bowshock_rho_flags!=0])) + return bowshock_rho_SDF + + else: # bowshock from user-defined variable and values + bowshock_dict_flags = make_region_flags(variable_dict, own_condition_dict, flag_type="01") + __, dict_sdf = vtkDelaunay3d_SDF(query_points, f.get_cell_coordinates(cellids[bowshock_dict_flags!=0])) + + return dict_sdf + + + +def RegionFlags(datafile, outfilen, regions=["all"], ignore_boundaries=True, region_flag_type="01", magnetopause_kwargs={}, region_conditions={}): + """Creates a sidecar .vlsv file with flagged cells for regions and boundaries in near-Earth plasma environment. + Region flags (named flag_region, flags are fractions of filled conditions or 1/0): magnetosheath, magnetosphere, cusps, lobe_N, lobe_S, central_plasma_sheet + Boundary signed distance flags (named "SDF_boundary", flags are signed distances to boundary in m with inside being negative distance): magnetopause, bowshock + + possible regions: "all", "boundaries" (magnetopause, bow shock), "large_areas" (boundaries + upstream, magnetosheath, magnetosphere), "magnetosphere", "bowshock", + "cusps", "lobes", "central_plasma_sheet" + + Note that different runs may need different tresholds for region parameters and region accuracy should be verified visually + Beta* convex hull is most likely the best magnetopause method for regions, for just magnetopause with different options use magnetopause.py + + :param datafile: .vlsv bulk file name (and path) + :param outfilen: sidecar .vlsv file name (and path) + :kword ignore_boundaries: True: do not take cells in the inner/outer boundaries of the simulation into account when looking for regions (applicable for cusps and CPS for now) + :kword region_flag_type: "01" or "fraction", whether flags are binary (all conditions must be satisfied) or fractions of how many given conditions are met + :kword magnetopause_method: default "beta_star", other options: see magnetopause.py, use alpha=None + :kword region_conditions: optional dict where keys are str region names and values are condition dictionaries, for setting own conditions for bow shock, cusps, lobes or central plasma sheet + """ + + if "all" in regions: + regions.extend(["magnetopause", "bowshock", "upstream", "magnetosheath", "magnetosphere", "cusps", "lobes", "central_plasma_sheet"]) + else: + if "large_areas" in regions: + regions.extend(["magnetopause", "bowshock", "upstream", "magnetosheath", "magnetosphere"]) + if "boundaries" in regions: + regions.extend(["magnetopause", "bowshock"]) + if "cusps" in regions or "central_plasma_sheet" in regions: # for cusps and central plasma sheet we need the magnetoshpere + regions.extend(["magnetosphere"]) + if "magnetosphere" in regions or "magnetosheath" in regions: + regions.extend(["magnetopause"]) + if "bowshock" in regions or "magnetosheath" in regions or "upstream" in regions: + regions.extend(["bowshock"]) + + + + f = pt.vlsvfile.VlsvReader(file_name=datafile) + cellids =f.read_variable("CellID") + [xmin, ymin, zmin, xmax, ymax, zmax] = f.get_spatial_mesh_extent() + all_points = f.get_cell_coordinates(cellids) + cellIDdict = f.get_cellid_locations() + + ## writer ## + writer = pt.vlsvfile.VlsvWriter(f, outfilen) + writer.copy_variables_list(f, ["CellID"]) + #writer.copy_variables(f, varlist=["proton/vg_rho" , "vg_beta_star", "vg_temperature", "vg_b_vol", "vg_J", "vg_beta", "vg_connection", "vg_boundarytype"]) + + + # variable data dictionary + variables = {} + + if f.check_variable("proton/vg_rho"): # vg-grid + + # dict {varible call name : variable vlsvreader name or varible call name : (variable vlsvreader name, operator)} + varnames = {"density": "proton/vg_rho", + "temperature": "vg_temperature", + "beta": "vg_beta", + "beta_star": "vg_beta_star", + "B_magnitude": ("vg_B_vol", "magnitude"), + "B_x": ("vg_B_vol", "x"), + "B_y": ("vg_B_vol", "y"), + "B_z": ("vg_B_vol", "z"), + "connection": "vg_connection", + "J_magnitude": ("vg_J", "magnitude")} + + boundaryname = "vg_boundarytype" + + + def errormsg(varstr): logging.warning("{} could not be read, will be ignored".format(varstr)) + + for varname, filevarname in varnames.items(): + if isinstance(filevarname, str): # no operator + try: + variables[varname] = f.read_variable(name=filevarname, cellids=-1) + except: + variables[varname] = np.full((len(cellids)), np.nan) + errormsg(filevarname) + else: + try: + variables[varname] = f.read_variable(name=filevarname[0], cellids=-1, operator=filevarname[1]) + except: + variables[varname] = np.full((len(cellids)), np.nan) + errormsg(filevarname) + + + + #connectivity_region = treshold_mask(variables["connection"], 0) + #betastar_region = treshold_mask(variables["beta_star"], [0.0, 0.5]) + #magnetosphere_proper = np.where((connectivity_region==1) | (betastar_region==1), 1, 0) + #write_flags(writer, magnetosphere_proper, 'flag_magnetosphere') + #return 0 + + + # upstream point # this could probably be done better than just a random point in sw + upstream_point = [xmax-10*R_E,0.0,0.0] + upstream_cellid = f.get_cellid(upstream_point) + upstream_index = cellIDdict[upstream_cellid] + + ## MAGNETOPAUSE ## + if "magnetopause" in regions: + if magnetopause_kwargs: + __, magnetopause_SDF = magnetopause.magnetopause(datafile, **magnetopause_kwargs) + else: + __, magnetopause_SDF = magnetopause.magnetopause(datafile, method="beta_star_with_connectivity", Delaunay_alpha=None) # default magnetopause: beta*+ B connectivity convex hull + write_flags(writer, magnetopause_SDF, 'SDF_magnetopause') + write_flags(writer, np.where(np.abs(magnetopause_SDF) < 5e6, 1, 0), "flag_magnetopause") + + # save some magnetopause values for later + magnetopause_density = np.mean(variables["density"][np.abs(magnetopause_SDF) < 5e6]) + #print(f"{magnetopause_density=}") + magnetopause_temperature = np.mean(variables["temperature"][np.abs(magnetopause_SDF) < 5e6]) + #print(f"{magnetopause_temperature=}") + + ## MAGNETOSPHERE ## + # magnetosphere from magnetopause SDF + if "magnetosphere" in regions: + magnetosphere = np.where(magnetopause_SDF<0, 1, 0) + write_flags(writer, magnetosphere, 'flag_magnetosphere') + + + ## BOW SHOCK ## #TODO: similar kwargs system as magnetopause? + # bow shock from rho + if "bowshock" in regions: + if "bowshock" in region_conditions: + bowshock = bowshock_SDF(f, variables, all_points, own_condition_dict=region_conditions["bowshock"]) + else: + bowshock = bowshock_SDF(f, variables, all_points) # default upstream rho method, might fail with foreshock + write_flags(writer, bowshock, 'SDF_bowshock') + write_flags(writer, np.where(np.abs(bowshock) < 5e6, 1, 0), "flag_bowshock") + + # magnetosphere+magnetosheath -area + inside_bowshock = np.where(bowshock<0, 1, 0) + + ## MAGNETOSHEATH ## + if "magnetosheath" in regions: + # magnetosheath from bow shock-magnetosphere difference + magnetosheath_flags = np.where((inside_bowshock & 1-magnetosphere), 1, 0) + write_flags(writer, magnetosheath_flags, 'flag_magnetosheath') + + # save magnetosheath density and temperature for further use + #magnetosheath_density = np.mean(variables["density"][magnetosheath_flags == 1]) + #print(f"{magnetosheath_density=}") + #magnetosheath_temperature = np.mean(variables["temperature"][magnetosheath_flags == 1]) + #print(f"{magnetosheath_temperature=}") + + ## UPSTREAM ## + if "upstream" in regions: + # upstream from !bowshock + write_flags(writer, 1-inside_bowshock, 'flag_upstream') + #write_flags(writer, inside_bowshock, 'flag_inside_bowshock') + + + + ## INNER MAGNETOSPHERE REGIONS ## + + if "magnetosphere" in regions: + if ignore_boundaries: + noBoundaries = np.where(f.read_variable(name=boundaryname, cellids=-1) == 1, 1, 0) # boundarytype 1: not a boundary + #mask_inBowshock= np.where(((inside_bowshock == 1) & (noBoundaries == 1)), 1, 0).astype(bool) # only search inner regions from inside the magnetosheath and magnetosphere + mask_inMagnetosphere = np.where(((magnetosphere == 1) & (noBoundaries == 1)), 1, 0).astype(bool) # + else: + #mask_inBowshock = inside_bowshock.astype(bool) # only search inner regions from inside the magnetosheath and magnetosphere + mask_inMagnetosphere = magnetosphere.astype(bool) # + + #print("upstream B:", variables["B_magnitude"][upstream_index]) + + # cusps + if "cusps" in regions: + if "cusps" in region_conditions: + cusp_conditions = region_conditions["cusps"] + else: + cusp_conditions = {"density": [variables["density"][upstream_index], None], + #"beta_star": [0.1, None], + "connection": [0.0, 2.5], # either closed, or open-closed/closed-open + "B_magnitude":[2*variables["B_magnitude"][upstream_index], None], + "J_magnitude": [variables["J_magnitude"][upstream_index], None] + } + + cusp_flags = make_region_flags(variables, cusp_conditions, flag_type=region_flag_type, mask=mask_inMagnetosphere) + write_flags(writer, cusp_flags, 'flag_cusps', mask_inMagnetosphere) + + + # magnetotail lobes + if "lobes" in regions: + if "lobes" in region_conditions: + lobes_conditions = region_conditions["lobes"] + else: + lobes_conditions = {"beta": [None, 0.1], # Koskinen: 0.003 + "connection": [0.5, 2.5], + "density": [None, variables["density"][upstream_index]], # Koskinen: 1e-8 + "temperature": [None, 3.5e6], # Koskinen: 3.5e6 K + } + + # lobes slightly other way + lobe_N_conditions = {"beta": [None, 0.1], + "connection": [0.5, 2.5], + "B_x":[0, None], + "B_magnitude":[None, 10*variables["B_magnitude"][upstream_index]] + } + lobe_N_flags = make_region_flags(variables, lobe_N_conditions, flag_type=region_flag_type) + + + lobe_S_conditions = {"beta": [None, 0.1], + "connection": [0.5, 2.5], + "B_x":[None, 0], + "B_magnitude":[None, 10*variables["B_magnitude"][upstream_index]] + } + lobe_S_flags = make_region_flags(variables, lobe_S_conditions, flag_type=region_flag_type) + + write_flags(writer, lobe_N_flags, 'flag_lobe_N') + write_flags(writer, lobe_S_flags, 'flag_lobe_S') + + lobes_flags = make_region_flags(variables, lobes_conditions, flag_type=region_flag_type) + write_flags(writer, lobes_flags, 'flag_lobes') + + + # lobe density from median densities? + #lobes_mask = np.where((lobes_flags > 0.9), 1, 0).astype(bool) + #lobe_density = np.mean(variables["density"][lobes_mask])#[lobes_allcells>0.9]) + #print(f"{lobe_density=}") + + # Central plasma sheet + if "central_plasma_sheet" in regions: + if "central_plasma_sheet" in region_conditions: + central_plasma_sheet_conditions = region_conditions["central_plasma_sheet"] + else: + central_plasma_sheet_conditions = {"density": [None, 1e6], # Wolf intro ?should not be but seems to work + "beta": [1.0, None], # Koskinen: 6 + "temperature": [2e6, None], # 5e7 from Koskinen + "J_magnitude": [1e-9, None], + } + + central_plasma_sheet_flags = make_region_flags(variables, central_plasma_sheet_conditions,flag_type=region_flag_type, mask=mask_inMagnetosphere) + write_flags(writer, central_plasma_sheet_flags, 'flag_central_plasma_sheet', mask_inMagnetosphere) + + + ## Other boundary layers, PSBL sometimes works + + # Plasma sheet boundary layer (PSBL) + # if "PSBL" in regions or "all" in regions: + #PSBL_conditions = {#"density": (1e7, 1.0), # Koskinen: 1e5 + # "density": [None, 1e7], + #"temperature": [0.5e6, 1e6],#(1e7, 2.0), # from Koskinen + #"temperature": [2e6, None], + # "beta": [0.1, 1.0], + #"B_magnitude":[10*variables["B_magnitude"][upstream_index], None] + # "J_magnitude": [2*variables["J_magnitude"][upstream_index], None], + # } + +def main(): + + datafile = "/wrk-vakka/group/spacephysics/vlasiator/3D/EGE/bulk/bulk.0002000.vlsv" + outfilen = "EGE_regions_t2000.vlsv" + + RegionFlags(datafile, outfilen, regions=["all"]) + + +if __name__ == "__main__": + + main() diff --git a/testpackage/run_compare.sh b/testpackage/run_compare.sh index 981445fe..244d1c4f 100755 --- a/testpackage/run_compare.sh +++ b/testpackage/run_compare.sh @@ -1,7 +1,7 @@ #!/bin/bash -l #SBATCH -t 00:30:00 #SBATCH -J analysator_testpackage_compare -#SBATCH --constraint="ukko|carrington" +#SBATCH --constraint="carrington" #SBATCH -p short #SBATCH -n 1 #SBATCH --array=1-10