Commit a460eb1f authored by Simeon Doetsch's avatar Simeon Doetsch

Moved repository to Gitlab

parent b002d27e
__pycache__/
/*.egg-info/
.ipynb_checkpoints
*.pyc
build/
dist/
.vscode/
test.ipynb
# plutoplot - Plutocode Plotting library
This is a Python 3 library for plotting output data from the [PLUTO code](http://plutocode.ph.unito.it/)
This library is under development and mostly specialized on the functionality I personally need. Nonetheless I will try to make the code as widely applicable and documented as possible.
## Requirements
```
Python >=3.6
numpy
matplotlib
```
## Installation
Install it with pip:
```
pip install git+https://gitlab.mpcdf.mpg.de/sdoetsch/plutoplot.git
```
Or download repository (as archive or with git)
```
python setup.py install
```
If you downloaded via git and want to keep updating via git, you can also use
```
python setup.py develop
```
This way the package is always read in from the git directory.
## Example
```python
import plutoplot as pp
sim = pp.Simulation("path_to_simulation")
print(sim.grid)
# Access to grid coordinates
sim.x1
sim.vx1
# Access to simulation steps
sim[3].rho
# time evolution of mean pressure
mean_prs = sim.reduce(lambda D : D.prs.mean())
```
## Concepts
`plutoplot` offers three main classes for handling simulation data:
- `Simulation`: for a simulation
- `PlutoData`: for a single simulation data frame
- `Grid`: for the PLUTO domain grid
For data loading the user only has to instantiate a `Simulation`, the grid
and the `PlutoData` objects are created from the `Simulation` when needed.
`PlutoData` uses lazy loading for the actual data, which means the data is
loaded when is first needed, not on object instantiation.
Each variable is loaded seperately (independent of PLUTO save format),
e.g. when only density is needed, pressure is never put into memory.
## Simulation instantiation
A `Simulation` can be instantiated with:
```python
Simulation(sim_dir='', format=None, coordinates=None)
```
- `sim_dir`: Simulation directory (directory with `pluto.ini` and `definitions.h`).
plutoplot searchs for the gridfile and simulation data first in `sim_dir`,
then in `sim_dir/data`, and then looks up the data directory in `pluto.ini`.
Default: Current working directory
- `format`: file format of the simulation data, currently supports `dbl`, `flt`, `vtk`
in both `single_file` and `multiple_files` mode.
Default: `dbl`, `flt`, `vtk`, are tried in that order
- `coordinates`: coordinate system of the simulation grid.
Supports `cartesian`, `spherical`, `polar`, `cylindrical`.
Only necessary for projecting the grind into a cartesian system (e.g. for plotting).
Default: Read coordinate system from `definitions.h`, using `cartesian` as fallback
## Access `pluto.ini` and `definitions.h`
`plutoplot` can load the PLUTO configuration files:
```python
sim = Simulation('sim_dir')
sim.ini['section']['option']
sim.definitions['options']
```
The returned objects for `pluto.ini` (`Pluto_ini`) and `definitions.h` (`Definitions_H`)
are thin wrappers around `OrderedDict`s.
## Data access
The simulation steps can be optained from the `Simulation` object with the subscript syntax:
```python
sim = Simulation('sim_dir')
initial = sim[0]
last = sim[-1]
```
It supports the Python conventions for indexing (zero indexed, negative numbers
are handled as `len(sim)-i`)
Which variables are saved by PLUTO are accesible with
```python
sim.vars
# or
sim[i].vars
# e.g. ['rho', 'vx1', 'vx2', 'vx3', 'prs']
```
The variable names in this list can be then accessed from the `PlutoData` objects:
```python
sim[0].rho
sim[-1].vx1
```
The data is then returned as Numpy arrays, and can be used as usual.
## Plotting
# [Plutoplot repository moved to Gitlab](https://gitlab.com/simske/plutoplot)
__version__ = '0.1'
from .plutodata import PlutoData
from .simulation import Simulation
from . import misc
import multiprocessing
import os
import subprocess
import numpy as np
import matplotlib.pyplot as plt
# local imports
from .plutodata import PlutoData
from .simulation import Simulation
def parameter_generator(sim: Simulation, plot_func, output_path,
plot_args: dict={'vmin': None, 'vmax': None},
save_args: dict={'bbox_inches': 'tight'}):
for i in range(sim.n):
yield (sim, i, plot_func, output_path, plot_args, save_args)
def generate_frame(args):
sim, i, plot_func, output_path, plot_args, save_args = args
fig = plot_func(sim[i], **plot_args)
fig.savefig("{}{:04d}.png".format(output_path, i), **save_args)
plt.close(fig)
del sim[i]
def render_frames_parallel(sim: Simulation, plot_func, output_path: str='',
plot_args: dict={'vmin': None, 'vmax': None},
save_args: dict={'bbox_inches': 'tight'}, verbose=True):
with multiprocessing.Pool() as p:
if verbose:
total = len(sim)
print(f"Rendering frame 0/{total} (0%)", end='')
for i,_ in enumerate(p.imap_unordered(generate_frame, parameter_generator(sim, plot_func, output_path,
plot_args, save_args))):
print(f"\rRendering frame {i}/{total} ({i/total*100:.1f}%)", end='')
else:
p.starmap(generate_frame, parameter_generator(sim, plot_func, output_path,
plot_args, save_args))
def generate_animation(sim: Simulation, plot_func, output_name: str='animation.mp4',
framerate: int=25,
plot_args: dict={'vmin': None, 'vmax': None},
save_args: dict={'bbox_inches': 'tight'}, verbose=True):
os.mkdir('tmp')
render_frames_parallel(sim, plot_func, 'tmp/', plot_args, save_args, verbose=True)
subprocess.run(['ffmpeg', '-f', 'lavfi', '-i', 'anullsrc=stereo',
'-framerate', '{:d}'.format(framerate), '-i', 'tmp/%04d.png',
'-shortest', '-c:v', 'libx264', '-pix_fmt', 'yuv420p',
'-c:a', 'aac', output_name])
subprocess.run(['rm', '-r', 'tmp'])
import numpy as np
def generate_coord_mapping(coordinates: str) -> dict:
mappings = {
'cartesian': {
'x': 'x1',
'y': 'x2',
'z': 'x3',
},
'polar': {
'r': 'x1',
'phi': 'x2',
'z': 'x2'
},
'cylindrical': {
'r': 'x1',
'z': 'x3'
},
'spherical': {
'r': 'x1',
'theta': 'x2',
'phi': 'x3',
}
}
if coordinates not in mappings:
raise NotImplementedError("Coordinate system {} not implemented".format(coordinates))
mapping = mappings[coordinates]
velocities = {}
for key, value in mapping.items():
velocities["v"+key] = "v"+value
mapping.update(velocities)
return mapping
def generate_tex_mapping(coordinates: str) -> dict:
mappings = {
'cartesian': {
'x1': 'x',
'x2': 'y',
'x3': 'z'
},
'cylindrical': {
'x1': 'r',
'x2': 'z'
},
'polar': {
'x1': 'r',
'x2': r'\phi',
'x3': 'z'
},
'spherical': {
'x1': 'r',
'x2': r'\theta',
'x3': r'\phi'
}
}
if coordinates not in mappings:
raise NotImplementedError("Tex mappings for {} not implemented".format(coordinates))
mapping = mappings[coordinates]
velocities = {}
for key, value in mapping:
velocities["v"+key] = "v_" + value
mapping.update(velocities)
mapping['rho'] = r'\rho'
mapping['prs'] = 'p'
return mapping
def generate_coordinate_mesh(coordinates, x1, x2):
if coordinates == 'cartesian':
return x1, x2
elif coordinates == 'spherical':
r, theta = np.meshgrid(x1, x2)
x = r * np.sin(theta)
y = r * np.cos(theta)
return x, y
elif coordinates == 'cylindrical':
return x1, x2
elif coordinates == 'polar':
r, phi = np.meshgrid(x1, x2)
x = r * np.cos(phi)
y = r * np.sin(phi)
return x, y
import numpy as np
from collections import OrderedDict
from itertools import zip_longest
from os.path import join
from .coordinates import generate_coordinate_mesh
class Grid:
"""
dims: dimensions
shape: shape of arrays
size: total cells
"""
def __init__(self, gridfile):
self.coordinates = None
self.mappings = {}
self.mappings_tex = {}
self.read_gridfile(gridfile)
def set_coordinate_system(self, coordinates, mappings={}, mappings_tex={}):
self.coordinates = coordinates
self.mappings = mappings
self.mappings_tex = mappings_tex
def read_gridfile(self, gridfile_path) -> None:
x = []
dims = []
with open(gridfile_path, 'r') as gf:
# read all dimensions
while True:
# read line by line, stop if EOF
line = gf.readline()
if not line:
break
# ignore comments
if line[0] == '#':
continue
# find line with resolution in dimension
splitted = line.split()
if len(splitted) == 1:
dim = int(splitted[0])
dims.append(dim)
# read all data from dimension, moves file pointer
data = np.fromfile(gf, sep=' ', count=dim*3).reshape(-1, 3)
# calculate center of cell, and difference between cells
x.append((np.sum(data[:, 1:], axis=1)/2, data[:, 2] - data[:, 1]))
# save in grid datastructure
for i, xn in enumerate(x, start=1):
setattr(self, "x{}".format(i), xn[0])
setattr(self, "dx{}".format(i), xn[1])
self.dims = tuple(dims)
shape = []
if self.dims[0] > 1:
shape.append(self.dims[0])
if self.dims[1] > 1:
shape.append(self.dims[1])
if self.dims[2] > 1:
shape.append(self.dims[2])
self.data_shape = tuple(shape)
self.size = np.product(self.dims)
def mesh(self):
return generate_coordinate_mesh(self.coordinates, self.x1, self.x2)
def __getattr__(self, name):
try:
mappings = object.__getattribute__(self, 'mappings')
return object.__getattribute__(self, mappings[name])
except KeyError:
raise AttributeError("{} has no attribute '{}'".format(type(self), name))
def __str__(self):
return "PLUTO Grid, Dimensions {}".format(self.dims)
__repr__ = __str__
class SimulationMetadata:
def __init__(self, data_dir, format) -> None:
self.read_vars(join(data_dir, '{}.out'.format(format)), format)
# read VTK offsets in file
if format == 'vtk':
if self.file_mode == 'single':
self.vtk_offsets = vtk_offsets(join(data_dir, 'data.0000.vtk'))
else:
self.vtk_offsets = {}
for var in self.vars:
self.vtk_offsets.update(vtk_offsets(join(data_dir, '{}.0000.vtk'.format(var))))
def read_vars(self, path, format) -> None:
"""Read simulation step data and written variables"""
with open(path, 'r') as f:
lines = f.readlines()
self.n = len(lines)
# prepare arrays
self.t = np.empty(self.n, float)
self.dt = np.empty(self.n, float)
self.nstep = np.empty(self.n, int)
# information for all steps the same
file_mode, endianness, *self.vars = lines[0].split()[4:]
for i, line in enumerate(lines):
self.t[i], self.dt[i], self.nstep[i] = line.split()[1:4]
if file_mode == 'single_file':
self.file_mode = 'single'
elif file_mode == 'multiple_files':
self.file_mode = 'multiple'
self.charsize = 8 if format == 'dbl' else 4
endianness = '<' if endianness == 'little' else '>'
if format == 'vtk': endianness = '>' # VTK has always big endian
self.binformat = "{}f{}".format(endianness, self.charsize)
def vtk_offsets(filename) -> dict:
"""
Read positions of vars in VTK legacy file
"""
offsets = {}
with open(filename, 'rb') as f:
for l in f:
if not l or l == b'\n':
continue
split = l.split()
# skip coordinates (read in via gridfile)
if split[0] in [i + b'_COORDINATES' for i in [b"X", b"Y", b"Z",]]:
f.seek(int(split[1]) * 4 + 1, 1)
if split[0] == b'CELL_DATA':
bytesize = int(split[1]) * 4
# save position of variables
if split[0] == b'SCALARS':
var = split[1].decode()
f.readline() # skip "LOOKUP_TABLE"
offsets[var] = f.tell()
f.seek(bytesize, 1)
return offsets
class Pluto_ini(OrderedDict):
"""Parser for Plutocode initialization file pluto.ini"""
class Section(OrderedDict):
def __init__(self, name):
super().__init__()
self.name = name
def _align(self):
length = []
for key, value in self.items():
if isinstance(value, str):
length.append((len(key), len(value)))
else:
length.append((len(key), *[len(v) for v in value]))
return [max(i) for i in zip_longest(*length, fillvalue=0)]
def __str__(self):
colwidth = self._align()
out = "[{}]\n\n".format(self.name)
for key, value in self.items():
out += "{}{}\n".format(rpad([key], colwidth), lpad(value, colwidth[1:]))
return out
def __init__(self, path):
super().__init__()
self.path = path
self.parse()
def parse(self, txt=None):
if txt is None:
with open(self.path, 'r') as f:
lines = [l.strip() for l in f.readlines()]
else:
lines = [l.strip() for l in txt.split("\n")]
section = None
for line in lines:
if not line:
continue
elif line[0] == '[' and line[-1] == ']':
section = line[1:-1]
self[section] = self.Section(section)
else:
segments = line.split()
if len(segments) > 2:
self[section][segments[0]] = segments[1:]
else:
self[section][segments[0]] = segments[1]
def __str__(self):
out = ""
for section in self.values():
out += str(section) + "\n\n"
return out
def write(self, path=None):
if path is None:
path = self.path
with open(path, 'w') as f:
f.write(str(self))
def rpad(text, colwidth):
if isinstance(text, str):
text = [text]
out = ""
for txt, col in zip(text, colwidth):
pad = col - len(txt) + 1
out += txt + " " * pad
return out
def lpad(text, colwidth, spacer=2):
if isinstance(text, str):
text = [text]
out = ""
for txt, col in zip(text, colwidth):
pad = col - len(txt) + spacer
out += " " * pad + txt
return out
class Definitions_h(OrderedDict):
base_opt = ['physics', 'dimensions', 'components', 'geometry', 'body_force',
'forced_turb', 'cooling', 'reconstruction', 'time_stepping',
'dimensional_splitting', 'ntracer', 'user_def_parameters']
physics_dep = {'hd': ['eos', 'entropy_switch', 'thermal_conduction', 'viscosity', 'rotating_frame'],
'rhd': ['eos', 'entropy_switch'],
'mhd': ['eos', 'entropy_switch', 'divb_control', 'background_field',
'ambipolar_diffusion', 'resistivity', 'hall_mhd',
'thermal_conduction', 'viscosity', 'rotating_frame'],
'rmhd': ['eos', 'entropy_switch', 'divb_control', 'resistivity'],
'cr_transport': ['eos', 'anisotropy'],
'advection': []}
def __init__(self, path: str):
super().__init__()
self.path = path
self.parse()
def parse(self, txt: str=None) -> None:
with open(self.path, 'r') as f:
lines = [l.strip() for l in f.readlines()]
for line in lines:
if not line:
continue
segments = line.split()
if segments[0] == "#define":
self[segments[1].lower()] = segments[2].lower()
def __getitem__(self, key):
return super().__getitem__(key.lower())
def __setitem__(self, key, value):
super().__setitem__(key.lower(), value)
def cached_property(func):
def wrapper(self):
cached_name = "_" + func.__name__
try:
return getattr(self, cached_name)
except AttributeError:
setattr(self, cached_name, func(self))
return getattr(self, cached_name)
return property(wrapper)
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import ScalarFormatter
from mpl_toolkits.axes_grid1 import make_axes_locatable
from .io import Grid
def plot(data: np.ndarray, grid: Grid, ax=None, label: str=None, figsize=None,
cbar=True, vmin=None, vmax=None, cmap=None, projection: bool=True) -> None:
"""Simple colorplot for 2-dim data"""
if ax is None:
if figsize is None:
x_size = 6.4
y_size = x_size * grid.dims[1] / grid.dims[0] / 1.1
figsize = (x_size, y_size)
_, ax = plt.subplots(figsize=figsize)
if projection:
X, Y = grid.mesh()
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
else:
X, Y = grid.x1, grid.x2
ax.set_xlabel("${}$".format(grid.mappings_tex['x1']))
ax.set_ylabel("${}$".format(grid.mappings_tex['x2']))
im = ax.pcolormesh(X, Y, data.T, vmin=vmin, vmax=vmax, cmap=cmap)
ax.set_aspect(1)
if cbar:
formatter = ScalarFormatter()
formatter.set_powerlimits((-2,2))
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="10%", pad=0.05)
plt.colorbar(im, label=label, format=formatter, cax=cax)
return ax.figure, ax
import os
import numpy as np
from .io import Grid
from .plotting import plot
class PlutoData(object):
def __init__(self, n: int=-1, simulation=None):
"""
Read PLUTO output file
n: output step number. Default: -1, uses last picture
wdir: path to data directory
coordinates: 'cartesian', 'cylindrical', 'polar', or 'spherical', only for names
vars [dict]: variables for manual construction
simulation [Simulation]: parent Simulation, missing attributes are tried to fetch from there
"""
# construct object from simulation simulation
self.simulation = simulation
self.format = simulation.format
self.binformat = simulation.metadata.binformat
self.file_mode = simulation.metadata.file_mode
self.wdir, self.grid, self.vars = simulation.data_dir, simulation.grid, simulation.vars
self.n, self.t, self.dt, self.nstep = n, simulation.t[n], simulation.dt[n], simulation.nstep[n]
self.data = {}
def __getattr__(self, name):
"""Get grid/data attributes from corresponding dict, or load it"""
# normal attributes
# grid
getattribute = object.__getattribute__
grid = getattribute(self, 'grid')
try:
return getattr(grid, name)
except AttributeError:
pass
# data
data = getattribute(self, 'data')
try:
return data[name]
except KeyError:
if name in getattribute(self, 'vars'):
getattribute(self, '_load_var')(name)
return data[name]
# simulation
try:
return getattr(getattribute(self, 'simulation'), name)
except