{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# 3D AMRVAC"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We create a Magrittetorch model from a snapshot of 3D AMRVAC hydrodynamics simulation.\n",
"The hydro model was kindly provided by Jan Bolte.\n",
"Currently, the AMRVAC binary files can not yet be used directly to extract the snapshot data.\n",
"Hence we use vtk package to read the corresponding `.vtu` files."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setup"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Import the required functionalty."
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import magrittetorch.tools.radiativetransferutils as rtutils\n",
"import magrittetorch.tools.mesher as mesher\n",
"import magrittetorch.tools.setup as setup\n",
"from scipy.spatial import cKDTree\n",
"import torch\n",
"import vtk\n",
"from vtk.util.numpy_support import vtk_to_numpy\n",
"from magrittetorch.model.model import Model # Model class\n",
"from magrittetorch.model.geometry.geometry import GeometryType # GeometryType enum\n",
"from magrittetorch.model.geometry.boundary import BoundaryType # BoundaryType enum\n",
"# import magritte.core as magritte # Core functionality\n",
"import plons # Loading phantom data\n",
"import numpy as np # Data structures\n",
"import warnings # Hide warnings\n",
"warnings.filterwarnings('ignore') # especially for yt\n",
"import yt # 3D plotting\n",
"import os\n",
"\n",
"from tqdm import tqdm # Progress bars\n",
"from astropy import constants, units # Unit conversions\n",
"from scipy.spatial import Delaunay # Finding neighbors\n",
"from yt.funcs import mylog # To avoid yt output \n",
"mylog.setLevel(40) "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Define a working directory (you will have to change this; it must be an **absolute path**)."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"wdir = \"/lhome/thomasc/Magrittetorch-examples/AMRVAC_3D/\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Create the working directory."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"!mkdir -p $wdir"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Define file names."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"input_file = os.path.join(wdir, 'model_AMRVAC_3D.vtu' ) # AMRVAC snapshot\n",
"model_file = os.path.join(wdir, 'model_AMRVAC_3D.hdf5') # Resulting Magritte model\n",
"lamda_file = os.path.join(wdir, 'co.txt' ) # Line data file"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We use a snapshot and data file that can be downloaded with the following links."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"input_link = \"https://owncloud.ster.kuleuven.be/index.php/s/mqtDyDSMPm2TjmG/download\"\n",
"lamda_link = \"https://home.strw.leidenuniv.nl/~moldata/datafiles/co.dat\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Dowload the snapshot and the linedata (``%%capture`` is just used to suppress the output)."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"%%capture\n",
"!wget $input_link --output-document $input_file\n",
"!wget $lamda_link --output-document $lamda_file"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Extract data"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The script below extracts the required data from the snapshot `.vtu` file."
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 89/89 [00:00<00:00, 17914.05it/s]\n",
"100%|██████████| 297984/297984 [00:04<00:00, 66798.15it/s]\n",
"100%|██████████| 2139521/2139521 [00:10<00:00, 208348.85it/s]\n"
]
}
],
"source": [
"# Create a vtk reader to read the AMRVAC file and extract its contents.\n",
"reader = vtk.vtkXMLUnstructuredGridReader()\n",
"reader.SetFileName(input_file)\n",
"reader.Update()\n",
"\n",
"# Extract the grid output\n",
"grid = reader.GetOutput()\n",
"\n",
"# Extract the number of cells\n",
"ncells = grid.GetNumberOfCells() \n",
"\n",
"# Extract cell data\n",
"cellData = grid.GetCellData()\n",
"for i in tqdm(range(cellData.GetNumberOfArrays())):\n",
" array = cellData.GetArray(i)\n",
" if (array.GetName() == 'rho'):\n",
" rho = vtk_to_numpy(array)\n",
" if (array.GetName() == 'temperature'):\n",
" tmp = vtk_to_numpy(array) * units.K\n",
" if (array.GetName() == 'v1'):\n",
" v_x = vtk_to_numpy(array) * units.cm/units.s\n",
" if (array.GetName() == 'v2'):\n",
" v_y = vtk_to_numpy(array) * units.cm/units.s\n",
" if (array.GetName() == 'v3'):\n",
" v_z = vtk_to_numpy(array) * units.cm/units.s\n",
"\n",
"# Convert rho (total density) to abundances\n",
"nH2 = rho * 1.0e+6 * constants.N_A.si.value / 2.02 * units.m**-3\n",
"nCO = nH2 * 1.0e-4\n",
"\n",
"# Convenience arrays\n",
"zeros = np.zeros(ncells)\n",
"ones = np.ones (ncells)\n",
"\n",
"# Convert to fractions of the speed of light\n",
"velocity = np.stack((v_x, v_y, v_z), axis=1)\n",
"# Define turbulence at 150 m/s\n",
"trb = 150.0 * ones * units.m/units.s\n",
"\n",
"# Extract cell centres as Magritte points\n",
"centres = []\n",
"for c in tqdm(range(ncells)):\n",
" cell = grid.GetCell(c)\n",
" centre = np.zeros(3)\n",
" for i in range(8):\n",
" centre = centre + np.array(cell.GetPoints().GetPoint(i))\n",
" centre = 0.125 * centre\n",
" centres.append(centre)\n",
"centres = np.array(centres) * units.cm\n",
"position = centres\n",
"\n",
"# We assume that the geometry to be rectangular. We still need to define the boundary however.\n",
"# For adaptive mesh refinement models, we put a boundary box around it (to ensure that we have a convex boundary around the model; this is mainly important for imaging the model).\n",
"\n",
"delta_xyz = (np.max(centres, axis=0) - np.min(centres, axis=0))/2.0#half the width in each direction\n",
"center_xyz = (np.max(centres, axis=0) + np.min(centres, axis=0))/2.0\n",
"xyz_max = center_xyz + 1.001*delta_xyz\n",
"xyz_min = center_xyz - 1.001*delta_xyz\n",
"\n",
"# Add a boundary box around the model, putting the boundary points on a regular grid\n",
"box = mesher.create_cubic_uniform_hull(xyz_min, xyz_max, order=5)#the boundary box contains 2**order+1 points in each direction.\n",
"nbox = len(box)\n",
"\n",
"# Find closest points, in order to extrapolate the values towards the boundary box.\n",
"corresp_points = cKDTree(position).query(box)[1]\n",
"\n",
"# Map data of model to box\n",
"position_box = box\n",
"velocity_box = velocity[corresp_points]\n",
"nCO_box = nCO [corresp_points]\n",
"nH2_box = nH2 [corresp_points]\n",
"tmp_box = tmp [corresp_points]\n",
"trb_box = trb [corresp_points]\n",
"\n",
"#and extend data arrays\n",
"nCO = np.concatenate((nCO_box, nCO))\n",
"nH2 = np.concatenate((nH2_box, nH2))\n",
"zeros = np.concatenate((np.zeros(nbox), zeros))\n",
"ones = np.concatenate((np.ones(nbox), ones))\n",
"velocity = np.concatenate((velocity_box, velocity), axis=0)\n",
"tmp = np.concatenate((tmp_box, tmp))\n",
"trb = np.concatenate((trb_box, trb))\n",
"position = np.concatenate((box, position), axis=0)\n",
"\n",
"\n",
"# Extract Delaunay vertices (= Voronoi neighbors)\n",
"#Note: if you know the neighbors structure of the mesh + boundary already, you can skip this step and directly use the neighbors structure, filling in the variables nbs and n_nbs.\n",
"delaunay = Delaunay(position)\n",
"(indptr, indices) = delaunay.vertex_neighbor_vertices\n",
"neighbors = [indices[indptr[k]:indptr[k+1]] for k in range(len(position))]\n",
"nbs = [n for sublist in neighbors for n in sublist]\n",
"n_nbs = [len(sublist) for sublist in neighbors]\n",
"\n",
"# Compute the indices of the boundary particles of the mesh, extracted from the Delaunay vertices\n",
"# Note: technically, we already know the boundary (the constructed box), but as we compute a Delaunay triangulation either way to find the neighbors, we might as well use it to compute the boundary points.\n",
"boundary = set([])\n",
"for i in tqdm(range(delaunay.neighbors.shape[0])):\n",
" for k in range(4):\n",
" if (delaunay.neighbors[i][k] == -1):\n",
" nk1,nk2,nk3 = (k+1)%4, (k+2)%4, (k+3)%4 \n",
" boundary.add(delaunay.simplices[i][nk1])\n",
" boundary.add(delaunay.simplices[i][nk2])\n",
" boundary.add(delaunay.simplices[i][nk3])\n",
"\n",
"boundary = list(boundary)\n",
"boundary = np.array(boundary)\n",
"#convert to torch tensor\n",
"boundary_torch = torch.from_numpy(boundary).type(torch.int64)\n",
"nbs_torch = torch.Tensor(nbs).type(torch.int64)\n",
"n_nbs_torch = torch.Tensor(n_nbs).type(torch.int64)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Create model"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now all data is read, we can use it to construct a Magrittetorch model.\n",
"\n",
"
\n",
"\n",
"Warning\n",
"\n",
"Including all radiative transitions can be computationally expensive (both in time and memory cost) for self-consistent NLTE radiative transfer. For LTE radiative transfer, this is not the case, altough if one wants to image a specific line, that line must be in the list of considered transitions. For these examples, we include the first 10 radiative transitions of CO (J=1-0 to J=10-9). To consider all transitions, use `setup.set_linedata_from_LAMDA_file` as in the commented line below it.\n",
"\n",
"
"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Not considering all radiative transitions on the data file but only the specified ones!\n",
"Writing model to: /lhome/thomasc/Magrittetorch-examples/AMRVAC_3D/model_AMRVAC_3D.hdf5\n"
]
}
],
"source": [
"model = Model(model_file) # Create model object\n",
"model.geometry.geometryType.set(GeometryType.General3D) # This is a 3D model\n",
"model.geometry.boundary.boundaryType.set(BoundaryType.AxisAlignedCube) # With a spherical boundary\n",
"\n",
"# In order to make unit conversions trivial, we use astropy quantities as input\n",
"model.geometry.points.position.set_astropy(position) # Set point positions\n",
"model.geometry.points.velocity.set_astropy(velocity) # Set point velocities\n",
"model.chemistry.species.abundance.set_astropy(np.stack([nCO, nH2, zeros/units.m**3], axis=1))# Set species number densities\n",
"model.chemistry.species.symbol.set(np.array(['CO', 'H2', 'e-'], dtype='S')) #Set species symbols; should correspond to the LAMDA file format\n",
"#Note: the dtype='S' is necessary to correctly save and read the species symbols to/from the hdf5 file\n",
"\n",
"model.thermodynamics.temperature.gas.set_astropy(tmp) # Set gas temperature\n",
"model.thermodynamics.turbulence.vturb.set_astropy(trb) # Set turbulence velocity\n",
"\n",
"# Set the neighbors\n",
"model.geometry.points.neighbors.set(nbs_torch) # Set neighbors\n",
"model.geometry.points.n_neighbors.set(n_nbs_torch) # Set number of neighbors\n",
"model.geometry.boundary.boundary2point.set(boundary_torch) # Set which points are boundary points\n",
"\n",
"# Conveniently, the remeshing function puts the boundary points in front of the positions array\n",
"model = setup.set_boundary_condition_CMB (model) # Set CMB as boundary condition\n",
"model = setup.set_uniform_rays (model, 12) # Number of rays for NLTE raytracing; has be of the form 12*2**n\n",
"\n",
"#As this example does not do NLTE, we might as well only consider the first 10 transitions of CO\n",
"model = setup.set_linedata_from_LAMDA_file(model, lamda_file, {'considered transitions': [i for i in range(10)]})\n",
"# model = setup.set_linedata_from_LAMDA_file(model, lamda_file) # Consider all transitions\n",
"model = setup.set_quadrature (model, 7) # Set number of frequency quadrature points for NLTE radiative transfer\n",
"\n",
"model.write()\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Plot model"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Load the data in a `yt` unstructured mesh."
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [],
"source": [
"ds = yt.load_unstructured_mesh(\n",
" connectivity = delaunay.simplices.astype(np.int64),\n",
" coordinates = position.to_value(units.cm), # yt expects cm not m \n",
" node_data = {('connect1', 'n'): nCO[delaunay.simplices].to_value(units.m**-3)}\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Plot a slice through the mesh orthogonal to the z-axis."
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
" "
],
"text/plain": [
""
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"sl = yt.SlicePlot (ds, 'z', ('connect1', 'n'))\n",
"sl.set_cmap (('connect1', 'n'), 'magma')\n",
"sl.zoom (1.2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Show mesh on the plot."
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
" "
],
"text/plain": [
""
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"sl = yt.SlicePlot (ds, 'z', ('connect1', 'n'))\n",
"sl.set_cmap (('connect1', 'n'), 'magma')\n",
"sl.zoom (1.0)\n",
"sl.annotate_mesh_lines (plot_args={'color':'lightgrey', 'linewidths':[.25]})"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In the [next example](2_reduce_AMRVAC_3D.ipynb) we demonstrate how to reduce this model as in [De Ceuster et al. (2020)](https://ui.adsabs.harvard.edu/abs/2020MNRAS.499.5194D/abstract)."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.4"
},
"widgets": {
"application/vnd.jupyter.widget-state+json": {
"state": {},
"version_major": 2,
"version_minor": 0
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}