{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Analytic spiral"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We create a Magrittetorch model from an analytic description of an Archimedean spiral outflow.\n",
"The description of the model is based on [Homan et al. (2015)](https://www.aanda.org/articles/aa/full_html/2015/07/aa25933-15/aa25933-15.html)."
]
},
{
"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.model.model as magritte\n",
"import magrittetorch.tools.mesher as mesher \n",
"import magrittetorch.tools.setup as setup\n",
"import numpy as np # Data structures\n",
"from astropy import units, constants # Unit conversions\n",
"from astropy.units import Quantity # For function arguments\n",
"import os\n",
"import torch\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 units, constants # Unit conversions\n",
"from scipy.spatial import Delaunay, cKDTree # Finding neighbors\n",
"from yt.funcs import mylog # To avoid yt output \n",
"mylog.setLevel(40) # as error messages"
]
},
{
"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/Analytic_spiral/\""
]
},
{
"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": [
"model_file = os.path.join(wdir, 'model_analytic_spiral.hdf5') # Resulting Magritte model\n",
"lamda_file = os.path.join(wdir, 'co.txt' ) # Line data file"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We use a data file that can be downloaded with the following links."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"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 $lamda_link --output-document $lamda_file"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Spiral parameters"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The functions below describe the spiral structure, based on the parameters in [Homan et al. (2015)](https://www.aanda.org/articles/aa/full_html/2015/07/aa25933-15/aa25933-15.html)."
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"phi_0 = 0 / 180.0 * np.pi * units.rad\n",
"alpha = 15 / 180.0 * np.pi * units.rad\n",
"sigma = 20 * constants.au\n",
"b = 270 / (2.0*np.pi) * constants.au\n",
"\n",
"\n",
"def r_0(r: Quantity[units.m], theta: Quantity[units.rad], phi: Quantity[units.rad]) -> Quantity[units.m]:\n",
" return b*(phi - phi_0)/units.rad\n",
"\n",
"\n",
"def sigma_r(r: Quantity[units.m], theta: Quantity[units.rad], phi: Quantity[units.rad]) -> Quantity[units.m]:\n",
" return sigma\n",
"\n",
"\n",
"def theta_0(r: Quantity[units.m], theta: Quantity[units.rad], phi: Quantity[units.rad]) -> Quantity[units.rad]:\n",
" return 0.5*np.pi * units.rad\n",
"\n",
"\n",
"def sigma_theta(r: Quantity[units.m], theta: Quantity[units.rad], phi: Quantity[units.rad]) -> Quantity[units.rad]:\n",
" return alpha\n",
"\n",
"\n",
"def S(r: Quantity[units.m], theta: Quantity[units.rad], phi: Quantity[units.rad]) -> Quantity[units.dimensionless_unscaled]:\n",
" \"\"\"\n",
" Spiral-shaped distribution function.\n",
" \"\"\"\n",
" rd = (r - r_0 (r,theta,phi)) / sigma_r (r,theta,phi)\n",
" td = (theta - theta_0(r,theta,phi)) / sigma_theta(r,theta,phi)\n",
" return np.exp(-0.5*(rd**2 + td**2))\n",
"\n",
"\n",
"def SS(r: Quantity[units.m], theta: Quantity[units.rad], phi: Quantity[units.rad]) -> Quantity[units.dimensionless_unscaled]:\n",
" \"\"\"\n",
" Spiral-shaped distribution function.\n",
" (Including 10 windings.)\n",
" \"\"\"\n",
" result = np.zeros_like(r) / units.m\n",
" for i in range(10):\n",
" result += S(r, theta, phi+2.0*np.pi*i*units.rad)\n",
" return result"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Physical parameters"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"R_max = 600.0 * constants.au\n",
"R_star = 1.0 * constants.R_sun\n",
"v_inf = 14.5e+3 * units.m/units.s\n",
"beta = 0.4 * units.dimensionless_unscaled\n",
"Mdot = 1.5e-6 * (constants.M_sun/units.year) # [kg/s]\n",
"R_0 = 2.0 * R_star # [m]\n",
"epsilon = 0.5 * units.dimensionless_unscaled\n",
"T_star = 2330.0 * units.K\n",
"T_0 = 60.0 * units.K\n",
"rho_max = 1.0e-11 *units.kg/units.m**3\n",
"\n",
"\n",
"def v_r(r: Quantity[units.m]) -> Quantity[units.m/units.s]:\n",
" \"\"\"\n",
" Beta-law velocity profile.\n",
" \"\"\"\n",
" return v_inf * (1.0 - R_0/r)**beta\n",
"\n",
"\n",
"def rho_regular(r: Quantity[units.m]) -> Quantity[units.kg/units.m**3]:\n",
" \"\"\"\n",
" Density profile assuming a regular, spherically symmetric outflow.\n",
" \"\"\"\n",
" return Mdot / (4.0*np.pi*r**2*v_r(r))\n",
"\n",
"\n",
"def rho_spiral(r: Quantity[units.m], theta: Quantity[units.rad], phi: Quantity[units.rad]) -> Quantity[units.kg/units.m**3]:\n",
" \"\"\"\n",
" Density profile assuming a spiral outflow.\n",
" \"\"\"\n",
" return rho_max * (2.0*np.pi*b/r)**2 * SS(r, theta, phi)\n",
"\n",
"\n",
"def rho_wind(r: Quantity[units.m], theta: Quantity[units.rad], phi: Quantity[units.rad]) -> Quantity[units.kg/units.m**3]:\n",
" \"\"\"\n",
" Wind density, combined regular and spiral outflow.\n",
" \"\"\"\n",
" return rho_regular(r) + rho_spiral(r, theta, phi)\n",
"\n",
"\n",
"def T_regular(r: Quantity[units.m]) -> Quantity[units.K]:\n",
" \"\"\"\n",
" Temperature profile assuming a regular, spherically symmetric outflow.\n",
" \"\"\"\n",
" return T_star * (R_star/r)**epsilon\n",
"\n",
"\n",
"def T_spiral(r: Quantity[units.m], theta: Quantity[units.rad], phi: Quantity[units.rad]) -> Quantity[units.K]:\n",
" \"\"\"\n",
" Temperature profile assuming a spiral outflow.\n",
" \"\"\"\n",
" return T_0 * SS(r, theta, phi)\n",
"\n",
"\n",
"def T_wind(r: Quantity[units.m], theta: Quantity[units.rad], phi: Quantity[units.rad]) -> Quantity[units.K]:\n",
" \"\"\"\n",
" Wind temperature, combined regular and spiral outflow.\n",
" \"\"\"\n",
" return T_regular(r) + T_spiral(r, theta, phi)"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [],
"source": [
"def spherical(x: Quantity[units.m],y: Quantity[units.m],z: Quantity[units.m])-> tuple[Quantity[units.m],Quantity[units.rad],Quantity[units.rad]]:\n",
" \"\"\"\n",
" Convert cartesian to spherical coordinates.\n",
" \"\"\"\n",
" r = np.sqrt (x**2 + y**2 + z**2)\n",
" t = np.arccos (z/r)\n",
" p = np.arctan2(y,x)\n",
" return (r,t,p)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Create background mesh (with desired mesh element sizes)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Define the desired mesh element size in a ($100 \\times 100 \\times 100$) cube."
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [],
"source": [
"resolution = 100\n",
"\n",
"xs = np.linspace(-R_max, +R_max, resolution, endpoint=True)\n",
"ys = np.linspace(-R_max, +R_max, resolution, endpoint=True)\n",
"zs = np.linspace(-R_max, +R_max, resolution, endpoint=True)\n",
"\n",
"(Xs, Ys, Zs) = np.meshgrid(xs, ys, zs)\n",
"\n",
"position = np.stack((Xs.ravel(), Ys.ravel(), Zs.ravel()), axis=1)\n",
"\n",
"# Compute the corresponding spherical coordinates\n",
"# Rs = np.sqrt (Xs**2 + Ys**2 + Zs**2)\n",
"# Ts = np.arccos (Zs/Rs)\n",
"# Ps = np.arctan2(Ys,Xs)\n",
"Rs, Ts, Ps = spherical(Xs,Ys,Zs)\n",
"\n",
"# Evaluate the density on the cube\n",
"rhos = rho_wind(Rs, Ts, Ps)\n",
"rhos = np.nan_to_num(rhos, nan=1.0*units.kg/units.m**3)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now we remesh the grid using the built-in remesher"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"new interior points: 583111\n",
"number boundary points: 1538\n"
]
}
],
"source": [
"positions_reduced, nb_boundary = mesher.remesh_point_cloud(position, rhos.ravel(), max_depth=10, threshold= 2e-1, hullorder=4)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We add a spherical inner boundary at 0.04*R_max"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"number of points in reduced grid: 584918\n"
]
}
],
"source": [
"healpy_order = 5 #using healpy to determine where to place the 12*healpy_order**2 boundary points on the sphere\n",
"origin = np.array([0.0,0.0,0.0]).T #the origin of the inner boundary sphere\n",
"positions_reduced, nb_boundary = mesher.point_cloud_add_spherical_inner_boundary(positions_reduced, nb_boundary, 0.04*R_max, healpy_order=healpy_order, origin=origin)\n",
"nb_boundary = 0 #inner boundary not yet implemented in magrittetorch, thus not classifying these first n boundary points as such\n",
"print(\"number of points in reduced grid: \", len(positions_reduced))\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We add a spherical outer boundary at R_max"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"number of points in reduced grid: 290761\n"
]
}
],
"source": [
"healpy_order = 15 #using healpy to determine where to place the 12*healpy_order**2 boundary points on the sphere\n",
"origin = np.array([0.0,0.0,0.0]).T #the origin of the outer boundary sphere\n",
"positions_reduced, nb_boundary = mesher.point_cloud_add_spherical_outer_boundary(positions_reduced, nb_boundary, R_max, healpy_order=healpy_order, origin=origin)\n",
"print(\"number of points in reduced grid: \", len(positions_reduced))"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [],
"source": [
"npoints = len(positions_reduced)\n",
"\n",
"XCO = 6.0e-4 * units.dimensionless_unscaled # [.]\n",
"vturb = 1.5e+3 * units.m/units.s # [m/s]\n",
"\n",
"\n",
"def density(rr: Quantity[units.m]) -> Quantity[units.kg/units.m**3]:\n",
" \"\"\"\n",
" Density function.\n",
" \"\"\"\n",
" (r, t, p) = spherical(rr[:, 0], rr[:, 1], rr[:, 2])\n",
" r[r < R_0] = R_0 #density not defined inside the star\n",
" return rho_wind(r, t, p)\n",
"\n",
"\n",
"def abn_nH2(rr: Quantity[units.m]) -> Quantity[units.m**-3]:\n",
" \"\"\"\n",
" H2 number density function.\n",
" \"\"\"\n",
" return density(rr) / (2.01588 * constants.u)\n",
"\n",
"\n",
"def abn_nCO(rr: Quantity[units.m]) -> Quantity[units.m**-3]:\n",
" \"\"\"\n",
" CO number density function.\n",
" \"\"\"\n",
" return XCO * abn_nH2(rr)\n",
"\n",
"\n",
"def temp(rr: Quantity[units.m]) -> Quantity[units.K]:\n",
" \"\"\"\n",
" Temperature function.\n",
" \"\"\"\n",
" (r, t, p) = spherical(rr[:, 0], rr[:, 1], rr[:, 2])\n",
" r[r < R_0] = R_0 #temperature not defined inside the star\n",
" return T_wind(r, t, p)\n",
"\n",
"\n",
"def velo(rr: Quantity[units.m]) -> Quantity[units.m/units.s]:\n",
" \"\"\"\n",
" Velocity function.\n",
" \"\"\"\n",
" (r, t, p) = spherical(rr[:, 0], rr[:, 1], rr[:, 2])\n",
" r[r < R_0] = R_0 #velocity not defined inside the star\n",
" return v_r(r)[:, None] * rr / r[:, None]\n",
"\n",
"\n",
"\n",
"# Extract Delaunay vertices (= Voronoi neighbors)\n",
"delaunay = Delaunay(positions_reduced)\n",
"(indptr, indices) = delaunay.vertex_neighbor_vertices\n",
"neighbors = [indices[indptr[k]:indptr[k+1]] for k in range(npoints)]\n",
"nbs = [n for sublist in neighbors for n in sublist]\n",
"n_nbs = [len(sublist) for sublist in neighbors]\n",
"\n",
"# Convenience arrays\n",
"zeros = np.zeros(npoints)\n",
"ones = np.ones (npoints)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Convert model functions to arrays based the model mesh."
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [],
"source": [
"position = positions_reduced\n",
"velocity = velo(positions_reduced)\n",
"nH2 = abn_nH2(positions_reduced)\n",
"nCO = abn_nCO(positions_reduced)\n",
"tmp = temp(positions_reduced)\n",
"trb = vturb * ones"
]
},
{
"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": 16,
"metadata": {
"tags": []
},
"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/Analytic_spiral/model_analytic_spiral.hdf5\n"
]
}
],
"source": [
"from magrittetorch.model.geometry.geometry import GeometryType\n",
"from magrittetorch.model.geometry.boundary import BoundaryType\n",
"\n",
"model = magritte.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.Sphere3D) # 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",
"model = setup.set_Delaunay_neighbor_lists (model) # Automatically computes and sets neighbors for each point, using a Delaunay triangulation\n",
"# For unitless quantities, we can also directly set the torch tensors\n",
"boundary_indices = torch.arange(nb_boundary, dtype = torch.int64)\n",
"model.geometry.boundary.boundary2point.set(boundary_indices) # Set which points are boundary points\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 transition 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 quadrature points\n",
"\n",
"model.write()"
]
},
{
"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": 17,
"metadata": {},
"outputs": [],
"source": [
"ds = yt.load_unstructured_mesh(\n",
" connectivity = delaunay.simplices.astype(np.int64),\n",
" coordinates = delaunay.points.astype(np.float64) * 100.0, # yt expects cm not m \n",
" node_data = {('connect1', 'n'): nCO[delaunay.simplices].astype(np.float64)}\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Plot a slice through the mesh orthogonal to the z-axis and x-axis."
]
},
{
"cell_type": "code",
"execution_count": 18,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
" "
],
"text/plain": [
""
]
},
"execution_count": 18,
"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": "code",
"execution_count": 19,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
" "
],
"text/plain": [
""
]
},
"execution_count": 19,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"sl = yt.SlicePlot (ds, 'x', ('connect1', 'n'))\n",
"sl.set_cmap (('connect1', 'n'), 'magma')\n",
"sl.zoom (1.2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Show meshes on the plots."
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
" "
],
"text/plain": [
""
]
},
"execution_count": 20,
"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)\n",
"sl.annotate_mesh_lines (plot_args={'color':'lightgrey', 'linewidths':[.25]})"
]
},
{
"cell_type": "code",
"execution_count": 21,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
" "
],
"text/plain": [
""
]
},
"execution_count": 21,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"sl = yt.SlicePlot (ds, 'x', ('connect1', 'n'))\n",
"sl.set_cmap (('connect1', 'n'), 'magma')\n",
"sl.zoom (1.2)\n",
"sl.annotate_mesh_lines (plot_args={'color':'lightgrey', 'linewidths':[.25]})"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"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
}