From 4e3290dc3fe96824faed34244dee61a711808d1a Mon Sep 17 00:00:00 2001 From: AlexHls Date: Thu, 20 Jan 2022 17:05:44 +0100 Subject: [PATCH] Added more profile options, CLI. Improved documentation [build docs] --- .../models/converters/arepo_to_tardis.ipynb | 149 ++++- docs/tardis.bib | 15 + tardis/io/parsers/arepo.py | 554 +++++++++++++++++- 3 files changed, 676 insertions(+), 42 deletions(-) diff --git a/docs/io/configuration/components/models/converters/arepo_to_tardis.ipynb b/docs/io/configuration/components/models/converters/arepo_to_tardis.ipynb index eac05f18c73..49c013283ae 100644 --- a/docs/io/configuration/components/models/converters/arepo_to_tardis.ipynb +++ b/docs/io/configuration/components/models/converters/arepo_to_tardis.ipynb @@ -8,6 +8,21 @@ "# Converting Arepo snapshots to be usable by TARDIS" ] }, + { + "cell_type": "markdown", + "id": "84aedaa4", + "metadata": {}, + "source": [ + "### What is [Arepo](https://arepo-code.org/)?\n", + "> Arepo is a massively parallel gravity and magnetohydrodynamics code for astrophysics, designed for problems of large dynamic range. It employs a finite-volume approach to discretize the equations of hydrodynamics on a moving Voronoi mesh, and a tree-particle-mesh method for gravitational interactions. Arepo is originally optimized for cosmological simulations of structure formation, but has also been used in many other applications in astrophysics.\n", + "\n", + "{cite:ps}`Weinberger2020`\n", + "\n", + "This parser is intended for loading Arepo output files ('snapshots'), extracting the relevant (line-of-sight dependent) data and exporting it to `csvy` files, which can in turn be used in TARDIS models ([How to run TARDIS with a custom model](../Custom_TARDIS_Model_Tutorial))\n", + "\n", + "*Note: This parser has been developed for the (not publically available) development version of Arepo, not the public version. Althought it should also work with snapshots from the public version, this has not been tested. If you run into trouble loading the snapshot using the built-in functions, try providing the data manually*" + ] + }, { "cell_type": "code", "execution_count": null, @@ -16,8 +31,18 @@ "outputs": [], "source": [ "import numpy as np\n", - "from tardis.io.parsers import arepo\n", - "import json" + "import matplotlib.pyplot as plt\n", + "import arepo\n", + "import json\n", + "%matplotlib inline" + ] + }, + { + "cell_type": "markdown", + "id": "60695b40", + "metadata": {}, + "source": [ + "### Loading the simulation data" ] }, { @@ -27,21 +52,22 @@ "source": [ "As a first step, the relevant data has to be loaded from an Arepo snapshot. In case you have the arepo-snap-util package installed, you can use the built in wrapper (as described below) to load the relevant data. In case you do not have this package installed or want to load the snapshot in a different way, you can manually provide the relevant data and continue with the next step.\n", "\n", - "If you're using the built-in tool, you will need to provide the path to the snapshot file, a list with the elements you want to include in your Tardis model, as well as the species file with which the Arepo simulation was run." + "If you're using the built-in tool, you will need to provide the path to the snapshot file, a list with the elements you want to include in your TARDIS model, as well as the species file with which the Arepo simulation was run. *(The species file should contain one header line, followed by two colums, where the first contains the names of all the species and is used to find the indices of the individual species within the snapshot. The second column is not used by the loader.)*" ] }, { - "cell_type": "code", - "execution_count": null, - "id": "6d97d9d3", + "cell_type": "markdown", + "id": "27d90938", "metadata": {}, - "outputs": [], "source": [ - "try:\n", - " snapshot = arepo.ArepoSnapshot(\"arepo_snapshot.hdf5\", [\"ni56\", \"si28\"], \"species55.txt\", resolution=32)\n", - " pos, vel, rho, xnuc, time = snapshot.get_grids()\n", - "except ImportError:\n", - " print(\"Looks like arepo-snap-util is missing!\")" + "In case you have the arepo-snap-util package installed, you can load the data directly from a snapshot:\n", + "```python\n", + "snapshot = arepo.ArepoSnapshot(\n", + " \"arepo_snapshot.hdf5\", [\"ni56\", \"si28\"], \"species55.txt\", resolution=32\n", + ")\n", + "pos, vel, rho, xnuc, time = snapshot.get_grids()\n", + "```\n", + "This will load the necessary from the snapshot. See the API documentation for more options on how to load snapshots." ] }, { @@ -49,9 +75,25 @@ "id": "49203e1f", "metadata": {}, "source": [ - "This will load the necessary from the snapshot. See the API documentation for more options on how to load snapshots.\n", + "This will fail with an error if you do not have the arepo-snap-util package installed. Since this is not a default dependency of TARDIS, lets manually load the data. *(This manual load can effectively be used to load all kinds of models unrelated to Arepo, as long as the data comes in the correct format.)*\n", "\n", - "As you can see, it will fail if you do not have the arepo-snap-util package installed. So lets manually load the data. In this case from a `json` file, where the data which would have been loaded with the function above has been manually saved to." + "In this case the data is loaded from a `json` file. This file has been created dumping the data which would have been loaded from the snapshot to a `json` file.\n", + "
\n", + " Code: (click to expand) \n", + " \n", + "```python\n", + "data = {\n", + " \"pos\" : pos.tolist(),\n", + " \"vel\" : vel.tolist(),\n", + " \"rho\" : rho.tolist(),\n", + " \"xnuc\": [xnuc[x].tolist() for x in list(xnuc.keys())],\n", + " \"time\": time,\n", + "}\n", + "json_string = json.dumps(data)\n", + "with open('arepo_snapshot.json', 'w') as outfile:\n", + " json.dump(json_string, outfile)\n", + "```\n", + "
" ] }, { @@ -103,10 +145,11 @@ "id": "bf82ae2e", "metadata": {}, "source": [ + "### Extracting a profile and converting it to a csvy file\n", "Now You can create a TARDIS model. There are three possibilities on how to extract the profiles from the snapshot:\n", - "- **Line profile**: This extracts the data along a straight line (the x-axis) without any averaging\n", - "- **Cone profile**: This extracts the data within a specified cone *(TBD)*\n", - "- **Full profile**: This averages over the whole simulation *(TBD)*" + "- **Line profile**: This extracts the data along a straight line (the x-axis) \n", + "- **Cone profile**: This extracts the data within a specified cone\n", + "- **Full profile**: This averages over the whole simulation" ] }, { @@ -116,7 +159,7 @@ "metadata": {}, "outputs": [], "source": [ - "profile = arepo.LineProfile(pos, vel, rho, xnuc, time)" + "profile = arepo.ConeProfile(pos, vel, rho, xnuc, time)" ] }, { @@ -125,7 +168,7 @@ "metadata": {}, "source": [ "This loads the data (here for the line profile), which can then be cut to the ranges which you want to include in your TARDIS model. The syntax for the other profiles is similar:\n", - "- `arepo.ConeProfile()`\n", + "- `arepo.LineProfile()`\n", "- `arepo.FullProfile()`" ] }, @@ -134,7 +177,8 @@ "id": "af9a76e0", "metadata": {}, "source": [ - "Next you can create the profiles acccording to the model option you selected. A diagnostic plot will be shown per default, but this behaviour can be turned off with the option `show_plot=False`. The plot will always show both the positve and negative axis." + "Next you can create the profiles acccording to the model option you selected. A diagnostic plot will be shown per default, but this behaviour can be turned off with the option `show_plot=False`. The plot will always show both the positve and negative axis.\n", + "> **_NOTE:_** The keyword `opening_angle=40` is only needed for the cone profile. The other modes do not accept this keyword! The angle itself is the opening angle of the full cone and NOT the angle between the central x-axis and the cone!" ] }, { @@ -144,7 +188,7 @@ "metadata": {}, "outputs": [], "source": [ - "profile.create_profile()" + "profile.create_profile(opening_angle=40)" ] }, { @@ -162,7 +206,7 @@ "metadata": {}, "outputs": [], "source": [ - "profile.create_profile(inner_radius=1e11, outer_radius=2e11)" + "profile.create_profile(opening_angle=40, inner_radius=1e11, outer_radius=2e11)" ] }, { @@ -170,7 +214,7 @@ "id": "dd6310ec", "metadata": {}, "source": [ - "Once you have created a profile of the desired region, you can export the profile to a `csvy` file, which in turn can be used in a TARDIS model. Here you have to specify how many shells you want to export. Note that the number of exported shells cannot be larger than the resolution of the given grid since no interpolation on the profile is done. If you want more shells, you either need to map your snapshot to a higher resolution Carthesian grid or interpolate the initial data yourself." + "Once you have created a profile of the desired region, you can export the profile to a `csvy` file, which in turn can be used in a TARDIS model. Here you have to specify how many shells you want to export. The profiles are rebinned using [Scipys binned_statistic function](https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.binned_statistic.html), using the mean value of the data in each bin." ] }, { @@ -180,7 +224,7 @@ "metadata": {}, "outputs": [], "source": [ - "profile.export(5, \"snapshot_converted_to_tardis.csvy\")" + "profile.export(50, \"snapshot_converted_to_tardis.csvy\")" ] }, { @@ -192,13 +236,64 @@ "\n", "All abundences will normalised such that roughly sum to 1, but slight deviations are expected to occur.\n", "\n", - "Note that the export function will not overwrite existing files, but rather add an incrementing number to the filename." + "> **_NOTE:_** The export function will not overwrite existing files, but rather add an incrementing number to the filename." + ] + }, + { + "cell_type": "markdown", + "id": "d0347042", + "metadata": {}, + "source": [ + "### Manually rebinning the data\n", + "Using `profile.rebin(, statistic=)`, you can manually rebin the data and use all `` keywords accepted by the `scipy.stats.binned_statistic` function. In this case you should pass the `statistic=None` keyword to the `export` function, so the data does not get rebinned twice." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9e9bcc38", + "metadata": {}, + "outputs": [], + "source": [ + "profile.create_profile(opening_angle=40, inner_radius=1e11, outer_radius=2e11)\n", + "profile.rebin(50)\n", + "profile.plot_profile()\n", + "plt.show()\n", + "profile.export(50, \"snapshot_converted_to_tardis.csvy\", statistic=None)" + ] + }, + { + "cell_type": "markdown", + "id": "8a4a41a8", + "metadata": {}, + "source": [ + "## Using the parser as a command line tool" + ] + }, + { + "cell_type": "markdown", + "id": "97bd4fa6", + "metadata": {}, + "source": [ + "You can also use the `tardis.io.arepo` package as a command line utility, e.g. if you are running batch jobs and want to automatically convert your snapshots from within you job-script.\n", + "\n", + "To export the same profile as in the example above you can run\n", + "```bash\n", + "python .//arepo.py snapshot.hdf5 snapshot_converted.csvy -o 40 --inner_radius 1e11 --outer_radius 2e11 -e ni56 si28 --save_plot plot.png --resolution 32 --plot_rebinned plot_binned.png\n", + "```\n", + "\n", + "This will also save diagnostic plots of both the raw and rebinned profiles. For more information on how to use the command line tool run\n", + "```bash\n", + "python .//arepo.py --help\n", + "```\n", + "\n", + "> **_NOTE:_** The command line tool does only work with snapshot files, not with e.g. json files. It is in any case not advised to " ] } ], "metadata": { "kernelspec": { - "display_name": "Python 3", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, @@ -212,7 +307,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.8" + "version": "3.7.12" } }, "nbformat": 4, diff --git a/docs/tardis.bib b/docs/tardis.bib index 14c30348b17..c56b1b0b164 100644 --- a/docs/tardis.bib +++ b/docs/tardis.bib @@ -270,3 +270,18 @@ @Article{Barna2017 doi = {10.1093/mnras/stx1894}, eprint = {1707.07848}, } + +@article{Weinberger2020, + title = {The AREPO Public Code Release}, + volume = {248}, + ISSN = {1538-4365}, + url = {http://dx.doi.org/10.3847/1538-4365/ab908c}, + DOI = {10.3847/1538-4365/ab908c}, + number = {2}, + journal = {The Astrophysical Journal Supplement Series}, + publisher = {American Astronomical Society}, + author = {Weinberger, Rainer and Springel, Volker and Pakmor, Rüdiger}, + year = {2020}, + month = {Jun}, + pages = {32} +} diff --git a/tardis/io/parsers/arepo.py b/tardis/io/parsers/arepo.py index f4a29afd624..04d42e2be71 100644 --- a/tardis/io/parsers/arepo.py +++ b/tardis/io/parsers/arepo.py @@ -6,6 +6,7 @@ import numpy as np import matplotlib.pyplot as plt import pandas as pd +from scipy import stats class ArepoSnapshot: @@ -78,7 +79,10 @@ def __init__( self.spec_ind = np.array(self.spec_ind) self.s = gadget_snap.gadget_snapshot( - filename, hdf5=True, quiet=True, lazy_load=True, + filename, + hdf5=True, + quiet=True, + lazy_load=True, ) rz_yaw = np.array( @@ -273,7 +277,7 @@ class Profile: """ def __init__(self, pos, vel, rho, xnuc, time): - """ + """ Parameters ----- pos : list of float @@ -395,12 +399,82 @@ def plot_profile(self, save=None, dpi=600, **kwargs): ) if save is not None: plt.savefig( - save, bbox_inches="tight", dpi=dpi, + save, + bbox_inches="tight", + dpi=dpi, ) return fig - def export(self, nshells, filename, direction="pos"): + def rebin(self, nshells, statistic="mean"): + """ + Rebins the data to nshells. Uses the scipy.stats.binned_statistic + to bin the data. The standard deviation of each bin can be obtained + by passing the statistics="std" keyword. + + Parameters + ----- + nshells : int + Number of bins of new data. + statistic : str + Scipy keyword for scipy.stats.binned_statistic. Default: mean + + Returns + ----- + self : Profile object + + """ + + self.vel_prof_p, bins_p = stats.binned_statistic( + self.pos_prof_p, + self.vel_prof_p, + statistic=statistic, + bins=nshells, + )[:2] + self.vel_prof_n, bins_n = stats.binned_statistic( + self.pos_prof_n, + self.vel_prof_n, + statistic=statistic, + bins=nshells, + )[:2] + + self.rho_prof_p = stats.binned_statistic( + self.pos_prof_p, + self.rho_prof_p, + statistic=statistic, + bins=nshells, + )[0] + self.rho_prof_n = stats.binned_statistic( + self.pos_prof_n, + self.rho_prof_n, + statistic=statistic, + bins=nshells, + )[0] + + for spec in self.species: + self.xnuc_prof_p[spec] = stats.binned_statistic( + self.pos_prof_p, + self.xnuc_prof_p[spec], + statistic=statistic, + bins=nshells, + )[0] + self.xnuc_prof_n[spec] = stats.binned_statistic( + self.pos_prof_n, + self.xnuc_prof_n[spec], + statistic=statistic, + bins=nshells, + )[0] + + self.pos_prof_p = np.array( + [(bins_p[i] + bins_p[i + 1]) / 2 for i in range(len(bins_p) - 1)] + ) + self.pos_prof_n = np.array( + [(bins_n[i] + bins_n[i + 1]) / 2 for i in range(len(bins_n) - 1)] + ) + + return self + + def export(self, nshells, filename, direction="pos", statistic="mean"): """ Function to export a profile as csvy file. Either the positive or negative direction can be exported. Does @@ -417,6 +491,9 @@ def export(self, nshells, filename, direction="pos"): Specifies if either the positive or negative direction is to be exported. Available options: ['pos', 'neg']. Default: pos + statistic : str + Scipy keyword for scipy.stats.binned_statistic. If + statistic=None, data is not rebinned. Default: "mean" Returns ----- @@ -474,7 +551,14 @@ def export(self, nshells, filename, direction="pos"): ) ) - f.write("".join(["\n", "---\n",])) + f.write( + "".join( + [ + "\n", + "---\n", + ] + ) + ) # WRITE DATA datastring = ["velocity,", "density,"] @@ -483,6 +567,10 @@ def export(self, nshells, filename, direction="pos"): datastring.append("%s" % self.species[-1].capitalize()) f.write("".join(datastring)) + # Rebin data to nshells + if statistic is not None: + self.rebin(nshells, statistic=statistic) + if direction == "pos": exp = [ self.vel_prof_p, @@ -500,12 +588,6 @@ def export(self, nshells, filename, direction="pos"): else: raise ValueError("Unrecognized option for keyword 'direction'") - # Selcect nshells equally spaced indices - if nshells > len(exp): - warnings.warn( - "nshells was grater then available resolution. Setting nshells to resolution" - ) - nshells = len(exp) inds = np.linspace(0, len(exp[0]) - 1, num=nshells, dtype=int) for i in inds: @@ -516,6 +598,19 @@ def export(self, nshells, filename, direction="pos"): return filename + def get_profiles(): + """Returns all profiles for manual post_processing etc.""" + return ( + self.pos_prof_p, + self.pos_prof_n, + self.vel_prof_p, + self.vel_prof_n, + self.rho_prof_p, + self.rho_prof_n, + self.xnuc_prof_p, + self.xnuc_prof_n, + ) + class LineProfile(Profile): """ @@ -531,8 +626,8 @@ def create_profile( save_plot=None, plot_dpi=600, ): - """ - Creates a profile along the x-axis without any averaging + """ + Creates a profile along the x-axis Parameters ----- @@ -658,8 +753,437 @@ def create_profile( class ConeProfile(Profile): - pass + """ + Class for profiles extracted inside a cone around the x-axis. + Extends Profile. + """ + + def create_profile( + self, + opening_angle=20.0, + inner_radius=None, + outer_radius=None, + show_plot=True, + save_plot=None, + plot_dpi=600, + ): + """ + Creates a profile along the x-axis without any averaging + + Parameters + ----- + opening_angle : float + Opening angle (in degrees) of the cone from which the + data is extracted. Refers to the total opening angle, not + the angle with respect to the x axis. Default: 20.0 + inner_radius : float + Inner radius where the profiles will be cut off. Default: None + outer_radius : float + Outer radius where the profiles will be cut off. Default: None + show_plot : bool + Specifies if a plot is to be shown after the creation of the + profile. Default: True + save_plot : str + Location where the plot is being saved. Default: None + plot_dpi : int + Dpi of the saved plot. Default: 600 + + Returns + ----- + profile : LineProfile object + + """ + + # Convert Carthesian coordinates into cylindrical coordinates + # P(x,y,z) -> P(x,r,theta) + cyl = np.array( + [ + self.pos[0], + np.sqrt(self.pos[1] ** 2 + self.pos[2] ** 2), + np.arctan(self.pos[2] / self.pos[1]), + ] + ) + + # Get maximum allowed r of points to still be in cone + dist = np.tan(opening_angle / 2) * np.abs(cyl[0]) + + # Create masks + cmask_p = np.logical_and(cyl[0] > 0, cyl[1] <= dist) + cmask_n = np.logical_and(cyl[0] < 0, cyl[1] <= dist) + + # Apply mask to data + pos_p = np.sqrt( + (self.pos[0][cmask_p]) ** 2 + + (self.pos[1][cmask_p]) ** 2 + + (self.pos[2][cmask_p]) ** 2 + ) + pos_n = np.sqrt( + self.pos[0][cmask_n] ** 2 + + self.pos[1][cmask_n] ** 2 + + self.pos[2][cmask_n] ** 2 + ) + + vel_p = np.sqrt( + self.vel[0][cmask_p] ** 2 + + self.vel[1][cmask_p] ** 2 + + self.vel[2][cmask_p] ** 2 + ) + vel_n = np.sqrt( + self.vel[0][cmask_n] ** 2 + + self.vel[1][cmask_n] ** 2 + + self.vel[2][cmask_n] ** 2 + ) + + rho_p = self.rho[cmask_p] + rho_n = self.rho[cmask_n] + + spec_p = {} + spec_n = {} + + for spec in self.species: + spec_p[spec] = self.xnuc[spec][cmask_p] + spec_n[spec] = self.xnuc[spec][cmask_n] + + self.pos_prof_p = np.sort(pos_p) + self.pos_prof_n = np.sort(pos_n) + + if outer_radius is None: + maxradius_p = max(self.pos_prof_p) + maxradius_n = max(self.pos_prof_n) + else: + maxradius_p = outer_radius + maxradius_n = outer_radius + + if inner_radius is None: + minradius_p = min(self.pos_prof_p) + minradius_n = min(self.pos_prof_n) + else: + minradius_p = inner_radius + minradius_n = inner_radius + + mask_p = np.logical_and( + self.pos_prof_p >= minradius_p, self.pos_prof_p <= maxradius_p + ) + mask_n = np.logical_and( + self.pos_prof_n >= minradius_n, self.pos_prof_n <= maxradius_n + ) + + if not mask_p.any() or not mask_n.any(): + raise ValueError("No points left between inner and outer radius.") + + self.rho_prof_p = np.array( + [x for _, x in sorted(zip(pos_p, rho_p), key=lambda pair: pair[0])] + )[mask_p] + self.rho_prof_n = np.array( + [x for _, x in sorted(zip(pos_n, rho_n), key=lambda pair: pair[0])] + )[mask_n] + + self.vel_prof_p = np.array( + [x for _, x in sorted(zip(pos_p, vel_p), key=lambda pair: pair[0])] + )[mask_p] + self.vel_prof_n = np.array( + [x for _, x in sorted(zip(pos_n, vel_n), key=lambda pair: pair[0])] + )[mask_n] + + for spec in self.species: + self.xnuc_prof_p[spec] = np.array( + [ + x + for _, x in sorted( + zip(pos_p, spec_p[spec]), key=lambda pair: pair[0] + ) + ] + )[mask_p] + self.xnuc_prof_n[spec] = np.array( + [ + x + for _, x in sorted( + zip(pos_n, spec_n[spec]), key=lambda pair: pair[0] + ) + ] + )[mask_n] + + self.pos_prof_p = self.pos_prof_p[mask_p] + self.pos_prof_n = self.pos_prof_n[mask_n] + + if show_plot: + self.plot_profile(save=save_plot, dpi=plot_dpi) + + return self class FullProfile(Profile): - pass + """ + Class for profiles extracted from the full snapshot, + i.e. angle averaged profiles. + Extends Profile. + """ + + def create_profile( + self, + inner_radius=None, + outer_radius=None, + show_plot=True, + save_plot=None, + plot_dpi=600, + ): + """ + Creates a profile from the full snapshot. Positive and negative + direction are identical. + + Parameters + ----- + inner_radius : float + Inner radius where the profiles will be cut off. Default: None + outer_radius : float + Outer radius where the profiles will be cut off. Default: None + show_plot : bool + Specifies if a plot is to be shown after the creation of the + profile. Default: True + save_plot : str + Location where the plot is being saved. Default: None + plot_dpi : int + Dpi of the saved plot. Default: 600 + + Returns + ----- + profile : LineProfile object + + """ + + pos_p = np.sqrt( + (self.pos[0]) ** 2 + (self.pos[1]) ** 2 + (self.pos[2]) ** 2 + ).flatten() + pos_n = np.sqrt( + self.pos[0] ** 2 + self.pos[1] ** 2 + self.pos[2] ** 2 + ).flatten() + + vel_p = np.sqrt( + self.vel[0] ** 2 + self.vel[1] ** 2 + self.vel[2] ** 2 + ).flatten() + vel_n = np.sqrt( + self.vel[0] ** 2 + self.vel[1] ** 2 + self.vel[2] ** 2 + ).flatten() + + rho_p = self.rho.flatten() + rho_n = self.rho.flatten() + + spec_p = {} + spec_n = {} + + for spec in self.species: + spec_p[spec] = self.xnuc[spec].flatten() + spec_n[spec] = self.xnuc[spec].flatten() + + self.pos_prof_p = np.sort(pos_p) + self.pos_prof_n = np.sort(pos_n) + + if outer_radius is None: + maxradius_p = max(self.pos_prof_p) + maxradius_n = max(self.pos_prof_n) + else: + maxradius_p = outer_radius + maxradius_n = outer_radius + + if inner_radius is None: + minradius_p = min(self.pos_prof_p) + minradius_n = min(self.pos_prof_n) + else: + minradius_p = inner_radius + minradius_n = inner_radius + + mask_p = np.logical_and( + self.pos_prof_p >= minradius_p, self.pos_prof_p <= maxradius_p + ) + mask_n = np.logical_and( + self.pos_prof_n >= minradius_n, self.pos_prof_n <= maxradius_n + ) + + if not mask_p.any() or not mask_n.any(): + raise ValueError("No points left between inner and outer radius.") + + self.rho_prof_p = np.array( + [x for _, x in sorted(zip(pos_p, rho_p), key=lambda pair: pair[0])] + )[mask_p] + self.rho_prof_n = np.array( + [x for _, x in sorted(zip(pos_n, rho_n), key=lambda pair: pair[0])] + )[mask_n] + + self.vel_prof_p = np.array( + [x for _, x in sorted(zip(pos_p, vel_p), key=lambda pair: pair[0])] + )[mask_p] + self.vel_prof_n = np.array( + [x for _, x in sorted(zip(pos_n, vel_n), key=lambda pair: pair[0])] + )[mask_n] + + for spec in self.species: + self.xnuc_prof_p[spec] = np.array( + [ + x + for _, x in sorted( + zip(pos_p, spec_p[spec]), key=lambda pair: pair[0] + ) + ] + )[mask_p] + self.xnuc_prof_n[spec] = np.array( + [ + x + for _, x in sorted( + zip(pos_n, spec_n[spec]), key=lambda pair: pair[0] + ) + ] + )[mask_n] + + self.pos_prof_p = self.pos_prof_p[mask_p] + self.pos_prof_n = self.pos_prof_n[mask_n] + + if show_plot: + self.plot_profile(save=save_plot, dpi=plot_dpi) + + return self + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "snapshot", + help="Snapshot file for which to create velocity profile plot", + ) + parser.add_argument( + "save", + help="Filename of exported .csvy file", + ) + parser.add_argument( + "-a", + "--alpha", + help="Euler angle alpha for rotation of desired direction to x-axis. Default: 0", + type=float, + default=0.0, + ) + parser.add_argument( + "-b", + "--beta", + help="Euler angle beta for rotation of desired direction to x-axis. Default: 0", + type=float, + default=0.0, + ) + parser.add_argument( + "-g", + "--gamma", + help="Euler angle gamma for rotation of desired direction to x-axis. Default: 0", + type=float, + default=0.0, + ) + parser.add_argument( + "-o", + "--opening_angle", + help="Opening angle of the cone from which profile is extracted. Default 20.0", + type=float, + default=20.0, + ) + parser.add_argument( + "-n", + "--nshells", + help="Number of shells to create. Default: 10", + type=int, + default=10, + ) + parser.add_argument( + "-x", + "--boxsize", + help="Size of the box (in cm) from which data is extracted. Default: 1e12", + type=float, + default=1e12, + ) + parser.add_argument( + "-e", + "--elements", + help="List of species to be included. Default: ni56", + default="ni56", + nargs="+", + ) + parser.add_argument( + "--eosspecies", + help="Species file including all the species used in the production of the composition file. Default: species55.txt", + default="species55.txt", + ) + parser.add_argument( + "--outer_radius", + help="Outer radius to which to build profile.", + type=float, + ) + parser.add_argument( + "--inner_radius", + help="Inner radius to which to build profile.", + type=float, + ) + parser.add_argument( + "--profile", + help="How to build profile. Available options: [line, cone, full]. Default: cone", + default="cone", + choices=["line", "cone", "full"], + ) + parser.add_argument( + "--resolution", + help="Resolution of Carthesian grid extracted from snapshot. Default: 512", + type=int, + default=512, + ) + parser.add_argument( + "--numthreads", + help="Number of threads used in snapshot tree walk. Default: 4", + type=int, + default=4, + ) + parser.add_argument("--save_plot", help="File name of saved plot.") + parser.add_argument( + "--dpi", help="Dpi of saved plot. Default: 600", type=int, default=600 + ) + parser.add_argument( + "--plot_rebinned", help="File name of plot after rebinning" + ) + + args = parser.parse_args() + + snapshot = ArepoSnapshot( + args.snapshot, + args.elements, + args.eosspecies, + alpha=args.alpha, + beta=args.beta, + gamma=args.gamma, + boxsize=args.boxsize, + resolution=args.resolution, + numthreads=args.numthreads, + ) + + pos, vel, rho, xnuc, time = snapshot.get_grids() + + if args.profile == "line": + profile = LineProfile(pos, vel, rho, xnuc, time) + elif args.profile == "cone": + profile = ConeProfile(pos, vel, rho, xnuc, time) + elif args.profile == "full": + profile = FullProfile(pos, vel, rho, xnuc, time) + + if args.profile == "cone": + profile.create_profile( + opening_angle=args.opening_angle, + inner_radius=args.inner_radius, + outer_radius=args.outer_radius, + save_plot=args.save_plot, + plot_dpi=args.dpi, + ) + else: + profile.create_profile( + inner_radius=args.inner_radius, + outer_radius=args.outer_radius, + save_plot=args.save_plot, + plot_dpi=args.dpi, + ) + + profile.export(args.nshells, args.save) + + if args.plot_rebinned: + profile.plot_profile(save=args.plot_rebinned, dpi=args.dpi)