# Copyright 2026 Capytaine developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tools to use xarray Datasets as inputs and outputs.
.. todo:: This module could be tidied up a bit and some methods merged or
uniformized.
"""
import logging
from datetime import datetime
from itertools import product
from collections import Counter
from typing import Sequence, List, Union, Optional
from pathlib import Path
import numpy as np
import pandas as pd
import xarray as xr
from capytaine import __version__
from capytaine.bodies.abstract_bodies import AbstractBody
from capytaine.bodies.bodies import FloatingBody
from capytaine.bodies.multibodies import Multibody
from capytaine.bodies.dofs import AbstractDof
from capytaine.bem.problems_and_results import (
LinearPotentialFlowProblem, DiffractionProblem, RadiationProblem,
LinearPotentialFlowResult, _default_parameters)
from capytaine.bem.airy_waves import airy_waves_pressure_dataarray
from capytaine.post_pro.kochin import compute_kochin
from capytaine.io.bemio import dataframe_from_bemio
LOG = logging.getLogger(__name__)
#########################
# Reading test matrix #
#########################
def _unsqueeze_dimensions(data_array, dimensions=None):
"""Add scalar coordinates as dimensions of size 1."""
if dimensions is None:
dimensions = list(data_array.coords.keys())
for dim in dimensions:
if len(data_array.coords[dim].values.shape) == 0:
data_array = xr.concat([data_array], dim=dim)
return data_array
[docs]
def problems_from_dataset(dataset: xr.Dataset,
body: AbstractBody,
) -> List[LinearPotentialFlowProblem]:
"""Generate a list of problems from a test matrix.
Parameters
----------
dataset : xarray Dataset
Test matrix containing the problems parameters.
body : AbstractBody
The body on which the computations of the test matrix will be applied.
Returns
-------
list of LinearPotentialFlowProblem
Raises
------
ValueError
if required fields are missing in the dataset
"""
if not isinstance(body, AbstractBody):
if isinstance(body, (list, tuple)) and isinstance(body[0], AbstractBody):
if len(body) == 1:
body = body[0]
else:
raise DeprecationWarning(
"Passing a list of bodies to `fill_dataset` or `problems_from_dataset` has been deprecated in version 3.\n"
"Consider using `problems_from_dataset` on each body separately and (if suitable) concatenating the results."
)
else:
raise TypeError(f"Unrecognized body type: {type(body)}. Expected a FloatingBody or a Multibody instance.")
# Should be done before looking for `frequency_keys`, otherwise
# frequencies provided as a scalar dimension will be skipped.
dataset = _unsqueeze_dimensions(dataset)
# SANITY CHECKS
# Warn user in case of key with unrecognized name (e.g. misspells)
keys_in_dataset = set(dataset.dims)
accepted_keys = {'wave_direction', 'radiating_dof', 'influenced_dof',
'omega', 'freq', 'period', 'wavelength', 'wavenumber',
'forward_speed', 'water_depth', 'rho', 'g', 'theta'}
unrecognized_keys = keys_in_dataset.difference(accepted_keys)
if len(unrecognized_keys) > 0:
LOG.warning(f"Unrecognized key(s) in dataset: {unrecognized_keys}")
if ("radiating_dof" not in keys_in_dataset) and ("wave_direction" not in keys_in_dataset):
raise ValueError("Neither 'radiating_dof' nor 'wave_direction' has been provided in the dataset. "
"No linear potential flow problem can be inferred.")
frequency_keys = keys_in_dataset & {'omega', 'freq', 'period', 'wavelength', 'wavenumber'}
if len(frequency_keys) > 1:
raise ValueError("Setting problems requires at most one of the following: omega (angular frequency) OR freq (in Hz) OR period OR wavenumber OR wavelength.\n"
"Received {}".format(frequency_keys))
# END SANITY CHECKS
if len(frequency_keys) == 0:
freq_type = "omega"
freq_range = [_default_parameters['omega']]
else: # len(frequency_keys) == 1
freq_type = list(frequency_keys)[0] # Get the only item
freq_range = dataset[freq_type].data
water_depth_range = dataset['water_depth'].data if 'water_depth' in dataset else [_default_parameters['water_depth']]
rho_range = dataset['rho'].data if 'rho' in dataset else [_default_parameters['rho']]
g_range = dataset['g'].data if 'g' in dataset else [_default_parameters['g']]
forward_speed_range = dataset['forward_speed'] if 'forward_speed' in dataset else [_default_parameters['forward_speed']]
wave_direction_range = dataset['wave_direction'].data if 'wave_direction' in dataset else None
radiating_dofs = dataset['radiating_dof'].data.astype(object) if 'radiating_dof' in dataset else None
# astype(object) is meant to convert Numpy internal string type numpy.str_ to Python general string type.
problems = []
if wave_direction_range is not None:
for freq, wave_direction, water_depth, forward_speed, rho, g \
in product(freq_range, wave_direction_range, water_depth_range,
forward_speed_range, rho_range, g_range):
problems.append(
DiffractionProblem(body=body, **{freq_type: freq},
wave_direction=wave_direction, water_depth=water_depth,
forward_speed=forward_speed, rho=rho, g=g)
)
if radiating_dofs is not None:
for freq, radiating_dof, water_depth, forward_speed, rho, g \
in product(freq_range, radiating_dofs, water_depth_range, forward_speed_range, rho_range, g_range):
if forward_speed == 0.0:
problems.append(
RadiationProblem(body=body, **{freq_type: freq},
radiating_dof=radiating_dof, water_depth=water_depth,
forward_speed=forward_speed, rho=rho, g=g)
)
else:
if wave_direction_range is None:
LOG.warning("Dataset contains non-zero forward speed (forward_speed=%.2f) but no wave_direction has been provided. Wave direction of 0 rad (x-axis) has been assumed.", forward_speed)
wave_direction_range = [0.0]
for wave_direction in wave_direction_range:
problems.append(
RadiationProblem(body=body, **{freq_type: freq},
radiating_dof=radiating_dof, water_depth=water_depth,
forward_speed=forward_speed, wave_direction=wave_direction,
rho=rho, g=g)
)
return sorted(problems)
########################
# Dataframe creation #
########################
def _detect_bemio_results(results, calling_function="_detect_bemio_results"):
error_msg = (
f"The function {calling_function} expected either a non-empty list of LinearPotentialFlowResult or a bemio.io object.\n"
f"Instead, it received:\n{repr(results)}"
)
if hasattr(results, '__iter__'):
if len(results) == 0:
raise ValueError("Iterable provided to `assemble_dataset` is empty.")
try:
if 'capytaine' in results[0].__module__:
bemio_import = False
else:
raise TypeError(error_msg)
except:
raise TypeError(error_msg)
else:
try:
if 'bemio.io' in results.__module__:
bemio_import = True
else:
raise TypeError(error_msg)
except:
raise TypeError(error_msg)
return bemio_import
[docs]
def assemble_dataframe(results, wavenumber=True, wavelength=True):
if _detect_bemio_results(results, calling_function="assemble_dataframe"):
return dataframe_from_bemio(results, wavenumber, wavelength) # TODO add hydrostatics
records_list = [record for result in results for record in result.records]
df = pd.DataFrame(records_list)
all_dofs_in_order = list({k: None for r in results for k in r.body.dofs.keys()})
# Using a dict above to remove duplicates while conserving ordering
dof_cat = pd.CategoricalDtype(categories=all_dofs_in_order)
df["influenced_dof"] = df["influenced_dof"].astype(dof_cat)
if 'added_mass' in df.columns:
df["radiating_dof"] = df["radiating_dof"].astype(dof_cat)
return df
######################
# Dataset creation #
######################
def _squeeze_dimensions(data_array, dimensions=None):
"""Remove dimensions if they are of size 1. The coordinates become scalar coordinates."""
if dimensions is None:
dimensions = data_array.dims
for dim in dimensions:
if len(data_array[dim]) == 1:
data_array = data_array.squeeze(dim, drop=False)
return data_array
def _dataset_from_dataframe(df: pd.DataFrame,
variables: Union[str, Sequence[str]],
dimensions: Sequence[str],
optional_dims: Sequence[str],
) -> Union[xr.DataArray, xr.Dataset]:
"""Transform a pandas.Dataframe into a xarray.Dataset.
Parameters
----------
df: pandas.DataFrame
the input dataframe
variables: string or sequence of strings
the variables that will be stored in the output dataset.
If a single name is provided, a DataArray of this variable will be provided instead.
dimensions: sequence of strings
Names of dimensions the variables depends on.
They will always appear as dimension in the output dataset.
optional_dims: sequence of strings
Names of dimensions the variables depends on.
They will appears as dimension in the output dataset only if they have
more than one different values.
"""
df = df.drop_duplicates(optional_dims + dimensions)
df = df.set_index(optional_dims + dimensions)
da = df.to_xarray()[variables]
da = _squeeze_dimensions(da, dimensions=optional_dims)
return da
def _rotation_center_data_array(body: AbstractBody) -> Optional[xr.DataArray]:
rotation_centers = np.array([b.rotation_center for b in body.bodies])
if not np.all(np.isnan(rotation_centers)):
hs = xr.DataArray(rotation_centers, dims=("body", "space_coordinate"))
hs.coords["space_coordinate"] = xr.DataArray(["x", "y", "z"], dims=["space_coordinate"])
hs.coords["body"] = xr.DataArray([b.name for b in body.bodies], dims=["body"])
hs = _squeeze_dimensions(hs, dimensions=["body"])
return hs
else:
LOG.debug("No rotation center could be found for any body. Skipping the `rotation_center` coordinate.")
return None
def _compute_hydrostatics_dataset(
body: AbstractBody,
*,
g: float = 9.81,
rho: float = 1000.0,
only_dofs: Optional[List[str]] = None,
) -> xr.Dataset:
hs = xr.Dataset()
hs.coords["space_coordinate"] = xr.DataArray(["x", "y", "z"], dims=["space_coordinate"])
hs.coords["g"] = xr.DataArray([g], dims=["g"])
hs.coords["rho"] = xr.DataArray([rho], dims=["rho"])
# Putting them first such that we have an helpful error message if the dofs or the center of mass is not defined
dof_cat = pd.CategoricalDtype(categories=list(body.dofs.keys()))
if only_dofs is None:
dofs = list(body.dofs.keys())
else:
if not all(d in body.dofs.keys() for d in only_dofs):
raise ValueError(f"Cannot keep only dofs {only_dofs} for body {body.__short_str__()} with dofs {list(body.dofs.keys())}")
dofs = only_dofs
hs.coords["radiating_dof"] = xr.DataArray(pd.Series(dofs, dtype=dof_cat).cat.remove_unused_categories(), dims=["radiating_dof"])
# Removing unused categories to match the behavior of `assemble_dataset` when the test matrix has only some of the possible radiating dofs.
hs.coords["influenced_dof"] = xr.DataArray(pd.Index(list(body.dofs.keys()), dtype=dof_cat), dims=["influenced_dof"])
hs.coords["influenced_dof"] = hs.coords["influenced_dof"].astype(str)
hs.coords["radiating_dof"] = hs.coords["radiating_dof"].astype(str)
if hasattr(body, 'hydrostatic_stiffness'):
hs["hydrostatic_stiffness"] = xr.DataArray([[body.hydrostatic_stiffness.sel(radiating_dof=dofs)]], dims=["g", "rho", "influenced_dof", "radiating_dof"])
else:
hs["hydrostatic_stiffness"] = xr.DataArray([[body.compute_hydrostatic_stiffness(rho=rho, g=g).sel(radiating_dof=dofs)]], dims=["g", "rho", "influenced_dof", "radiating_dof"])
if hasattr(body, 'inertia_matrix'):
hs["inertia_matrix"] = xr.DataArray([body.inertia_matrix.sel(radiating_dof=dofs)], dims=["rho", "influenced_dof", "radiating_dof"])
else:
hs["inertia_matrix"] = xr.DataArray([body.compute_rigid_body_inertia(rho=rho).sel(radiating_dof=dofs)], dims=["rho", "influenced_dof", "radiating_dof"])
# Other hydrostatics data
if isinstance(body, FloatingBody):
body = Multibody([body])
# Afterwards, treat even the single body as a one-element Multibody
# Not done before, because a single-body Multibody does not have the same dof naming than a FloatingBody.
hs.coords["body"] = xr.DataArray([b.name for b in body.bodies], dims=["body"])
hs["center_of_buoyancy"] = xr.DataArray(list(body.center_of_buoyancy.values()), dims=("body", "space_coordinate"))
hs["draught"] = xr.DataArray([np.abs(b.mesh.vertices[:, 2].min()) for b in body.bodies], dims=("body"))
hs["disp_mass"] = xr.DataArray(
[[b.disp_mass(rho=rho) for b in body.bodies]],
dims=("rho", "body"),
attrs=dict(long_name="Displaced mass", units="kg")
)
try:
hs.coords["center_of_mass"] = xr.DataArray(list(body.center_of_mass.values()), dims=("body", "space_coordinate"))
except Exception as e:
LOG.debug(f"No center of mass could be found for some body. Skipping the `center_of_mass` coordinate.\nError message: {e}")
return hs
[docs]
def compute_hydrostatics_dataset(
body: AbstractBody,
*,
g: Union[float, Sequence[float]] = 9.81,
rho: Union[float, Sequence[float]] = 1000.0,
only_dofs: Optional[List[str]] = None,
):
"""
Parameters
----------
body: AbstractBody
g: float or sequence of floats, optional
Gravitational acceleration(s) to consider. If a sequence is provided, the output dataset will contain one entry per value of g.
rho: float or sequence of floats, optional
Fluid density(ies) to consider. If a sequence is provided, the output dataset will contain one entry per value of rho.
only_dofs: list of strings, optional
If provided, only the radiating dofs in this list will be kept in the output dataset. By default, all dofs of the body are kept.
Meant to mimic the behavior of `fill_dataset` when the test matrix has only some of the possible radiating dofs.
Returns
-------
xarray.Dataset
"""
if isinstance(g, np.ndarray) and g.shape == ():
g = [float(g)]
elif isinstance(g, (int, float)):
g = [g]
if isinstance(rho, np.ndarray) and rho.shape == ():
rho = [float(rho)]
elif isinstance(rho, (int, float)):
rho = [rho]
datasets = []
for g_ in g:
for rho_ in rho:
datasets.append(_compute_hydrostatics_dataset(body, g=g_, rho=rho_, only_dofs=only_dofs))
hs = xr.merge(datasets, compat="no_conflicts", join="outer")
optional_dims = ["g", "rho", "body"]
hs = _squeeze_dimensions(hs, dimensions=optional_dims)
_rotation_center_da = _rotation_center_data_array(body)
if _rotation_center_da is not None:
hs.coords['rotation_center'] = _rotation_center_da
return hs
[docs]
def kochin_data_array(results: Sequence[LinearPotentialFlowResult],
theta_range: Sequence[float],
**kwargs,
) -> xr.Dataset:
"""Compute the Kochin function for a list of results and fills a dataset.
.. seealso::
:meth:`~capytaine.post_pro.kochin.compute_kochin`
The present function is just a wrapper around :code:`compute_kochin`.
"""
# TODO: this not very good to mix computation and data manipulation here...
records = pd.DataFrame([
dict(**result.problem._asdict(), theta=theta, kochin=kochin, kind=result.__class__.__name__)
for result in results
for theta, kochin in zip(theta_range.data,
compute_kochin(result, theta_range, **kwargs))
])
main_freq_type = Counter((res.provided_freq_type for res in results)).most_common(1)[0][0]
kochin_data = xr.Dataset()
if "RadiationResult" in set(records['kind']):
radiation = _dataset_from_dataframe(
records[records['kind'] == "RadiationResult"],
variables=['kochin'],
dimensions=[main_freq_type, 'radiating_dof', 'theta'],
optional_dims=['g', 'rho', 'water_depth', 'forward_speed', 'wave_direction']
)
kochin_data['kochin_radiation'] = radiation['kochin']
if "DiffractionResult" in set(records['kind']):
diffraction = _dataset_from_dataframe(
records[records['kind'] == "DiffractionResult"],
['kochin'],
dimensions=[main_freq_type, 'wave_direction', 'theta'],
optional_dims=['g', 'rho', 'water_depth', 'forward_speed']
)
kochin_data['kochin_diffraction'] = diffraction['kochin']
return kochin_data
VARIABLES_ATTRIBUTES = {
"omega": {
'long_name': 'Angular frequency',
'units': 'rad/s',
},
"freq": {
'long_name': 'Frequency',
'units': 'Hz',
},
"period": {
'long_name': 'Period',
'units': 's',
},
"wavenumber": {
'long_name': "Angular wavenumber",
'units': 'rad/m',
},
"wavelength": {
'long_name': "Wave length",
'units': 'm',
},
"encounter_omega": {
'long_name': "Encounter angular frequency",
'units': 'rad/s',
},
"encounter_wave_direction": {
'long_name': "Encounter wave direction",
'units': 'rad',
},
"wave_direction": {
'long_name': "Wave direction",
'units': "rad"
},
"radiating_dof": {
'long_name': 'Radiating DOF',
},
"influenced_dof": {
'long_name': 'Influenced DOF',
},
"added_mass": {
'long_name': 'Added mass',
},
"radiation_damping": {
'long_name': 'Radiation damping',
},
"diffraction_force": {
'long_name': "Diffraction force",
},
"Froude_Krylov_force": {
'long_name': "Froude Krylov force",
},
"pressure_radiation": {
'long_name': "Radiation pressure",
'units': 'Pa',
},
"pressure_diffraction": {
'long_name': "Diffraction pressure",
'units': 'Pa',
},
"pressure_Froude_Krylov": {
'long_name': "Froude Krylov pressure",
'units': 'Pa',
},
"sources_diffraction": {
'long_name': "Diffraction source distribution on the hull",
},
"sources_radiation": {
'long_name': "Radiation source distribution on the hull",
},
"sources_lid_diffraction": {
'long_name': "Diffraction source distribution on the lid",
},
"sources_lid_radiation": {
'long_name': "Radiation source distribution on the lid",
},
"potential_diffraction": {
'long_name': "Diffraction potential on the hull",
},
"potential_radiation": {
'long_name': "Radiation potential on the hull",
},
"potential_lid_diffraction": {
'long_name': "Diffraction potential on the lid",
},
"potential_lid_radiation": {
'long_name': "Radiation potential on the lid",
},
}
def _face_data_array(result, name, values, on_hull):
"""Part of a distribution defined on all the faces of the mesh including lid
(`values`, or None if missing) that is on the hull (or on the lid), as a
DataArray with the metadata of the problem.
The hull and lid are separated with the hull mask of the body, the order of
the faces in the mesh including lid being not necessarily hull first."""
mask = result.body.hull_mask
mask = mask if on_hull else ~mask
if values is None:
values = np.full((int(np.count_nonzero(mask)),), np.nan + 1j*np.nan)
else:
values = np.asarray(values)[mask]
dataarray = xr.DataArray(values, dims=["hull_face" if on_hull else "lid_face"], name=name)
return result.problem._wrap_dataarray(dataarray)
def _add_face_distributions(dataset, results, name, optional_dims, lid=True):
"""Add to the dataset the variables '<name>_diffraction' and '<name>_radiation'
(and '<name>_lid_diffraction', '<name>_lid_radiation' on the lid, if `lid`) for the distributions `name` of the results."""
if not any(getattr(r, name) is not None for r in results):
return
for kind, problem_class, dims in [
("diffraction", DiffractionProblem, optional_dims),
("radiation", RadiationProblem, optional_dims + ['wave_direction'])]:
selected = [r for r in results if isinstance(r.problem, problem_class)]
if len(selected) == 0:
continue
for on_hull, var_name in [(True, f"{name}_{kind}"), (False, f"{name}_lid_{kind}")]:
if not on_hull and (not lid or all(r.body.lid_mesh is None for r in selected)):
continue
dataset[var_name] = _squeeze_dimensions(xr.merge(
[_face_data_array(r, "distribution", getattr(r, name), on_hull) for r in selected],
compat="no_conflicts", join="outer"
)["distribution"], dimensions=dims)
[docs]
def assemble_dataset(results,
omega=True, freq=True, wavenumber=True, wavelength=True, period=True,
mesh=False, hydrostatics=True, attrs=None) -> xr.Dataset:
"""Transform a list of :class:`LinearPotentialFlowResult` into a :class:`xarray.Dataset`.
Parameters
----------
results: list of LinearPotentialFlowResult or BEMIO dataset
The results that will be read.
omega: bool, optional
If True, the coordinate 'omega' will be added to the output dataset.
freq: bool, optional
If True, the coordinate 'freq' will be added to the output dataset.
wavenumber: bool, optional
If True, the coordinate 'wavenumber' will be added to the output dataset.
wavelength: bool, optional
If True, the coordinate 'wavelength' will be added to the output dataset.
period: bool, optional
If True, the coordinate 'period' will be added to the output dataset.
mesh: bool, optional
If True, store the mesh (vertices and centers of each face of the hull,
and of the lid if any) and the dofs (motion, and when available its
gradient, of each dof evaluated on the hull) in the output dataset, as
the data variables ``mesh_vertices``, ``mesh_faces_center``,
``lid_mesh_vertices``, ``lid_mesh_faces_center``, ``dof_motions`` and
``dof_gradient_of_motions``.
hydrostatics: bool, optional
If True, store the hydrostatic data in the output dataset if they exist.
attrs: dict, optional
Attributes that should be added to the output dataset.
"""
bemio_import = _detect_bemio_results(results, calling_function="assemble_dataset")
if not bemio_import:
if len(results) == 0:
raise ValueError("No results provided to `assemble_dataset`. Please provide a non-empty list of `LinearPotentialFlowResult` or a bemio.io dataset.")
# We want to test that all results comes from the same body.
# Checking that all body objects are the same is tricky because clipping or parallelisation can cause some body to be copies.
# We only check that all bodies have a few properties in common.
bodies = {(r.body.name, r.body.mesh.nb_faces, r.body.nb_dofs) for r in results}
if len(bodies) > 1:
raise ValueError("Results from different bodies have been provided. This is not supported by `assemble_dataset`.\n"
"Consider using `assemble_dataset` on each body separately and concatenating the results.")
records = assemble_dataframe(results)
if bemio_import:
main_freq_type = "omega"
else:
main_freq_type = Counter((res.provided_freq_type for res in results)).most_common(1)[0][0]
if np.any(records["free_surface"] != 0.0):
LOG.warning("Datasets only support cases with a free surface (free_surface=0.0).\n"
"Cases without a free surface (free_surface=inf) are ignored.\n"
"See also https://github.com/mancellin/capytaine/issues/88")
records = records[records["free_surface"] == 0.0]
if attrs is None:
attrs = {}
attrs['creation_of_dataset'] = datetime.now().isoformat()
kinds_of_results = set(records['kind'])
optional_dims = ['g', 'rho', 'water_depth', 'forward_speed']
dataset = xr.Dataset()
if not bemio_import:
_rotation_center_da = _rotation_center_data_array(results[0].body)
if _rotation_center_da is not None:
dataset.coords['rotation_center'] = _rotation_center_da
# RADIATION RESULTS
if "RadiationResult" in kinds_of_results:
radiation_cases = _dataset_from_dataframe(
records[records['kind'] == "RadiationResult"],
variables=['added_mass', 'radiation_damping'],
dimensions=[main_freq_type, 'influenced_dof', 'radiating_dof'],
optional_dims=optional_dims + ['wave_direction'])
dataset = xr.merge([dataset, radiation_cases], compat="no_conflicts", join="outer")
# DIFFRACTION RESULTS
if "DiffractionResult" in kinds_of_results:
diffraction_cases = _dataset_from_dataframe(
records[records['kind'] == "DiffractionResult"],
variables=['diffraction_force', 'Froude_Krylov_force'],
dimensions=[main_freq_type, 'wave_direction', 'influenced_dof'],
optional_dims=optional_dims)
dataset = xr.merge([dataset, diffraction_cases], compat="no_conflicts", join="outer")
dataset['excitation_force'] = dataset['Froude_Krylov_force'] + dataset['diffraction_force']
# OTHER FREQUENCIES TYPES
if omega and main_freq_type != "omega":
omega_ds = _dataset_from_dataframe(
records,
variables=['omega'],
dimensions=[main_freq_type],
optional_dims=['g', 'water_depth'] if main_freq_type in {'wavelength', 'wavenumber'} else []
)
dataset.coords['omega'] = omega_ds['omega']
if freq and main_freq_type != "freq":
freq_ds = _dataset_from_dataframe(
records,
variables=['freq'],
dimensions=[main_freq_type],
optional_dims=['g', 'water_depth'] if main_freq_type in {'wavelength', 'wavenumber'} else []
)
dataset.coords['freq'] = freq_ds['freq']
if period and main_freq_type != "period":
period_ds = _dataset_from_dataframe(
records,
variables=['period'],
dimensions=[main_freq_type],
optional_dims=['g', 'water_depth'] if main_freq_type in {'wavelength', 'wavenumber'} else []
)
dataset.coords['period'] = period_ds['period']
if wavenumber and main_freq_type != "wavenumber":
wavenumber_ds = _dataset_from_dataframe(
records,
variables=['wavenumber'],
dimensions=[main_freq_type],
optional_dims=['g', 'water_depth'] if main_freq_type in {'period', 'omega'} else []
)
dataset.coords['wavenumber'] = wavenumber_ds['wavenumber']
if wavelength and main_freq_type != "wavelength":
wavelength_ds = _dataset_from_dataframe(
records,
variables=['wavelength'],
dimensions=[main_freq_type],
optional_dims=['g', 'water_depth'] if main_freq_type in {'period', 'omega'} else []
)
dataset.coords['wavelength'] = wavelength_ds['wavelength']
if not all(records["forward_speed"] == 0.0):
omegae_ds = _dataset_from_dataframe(
records,
variables=['encounter_omega'],
dimensions=['forward_speed', 'wave_direction', main_freq_type],
optional_dims=['g', 'water_depth'],
)
dataset.coords['encounter_omega'] = omegae_ds['encounter_omega']
encounter_wave_direction_ds = _dataset_from_dataframe(
records,
variables=['encounter_wave_direction'],
dimensions=['forward_speed', 'wave_direction', main_freq_type],
optional_dims=[],
)
dataset.coords['encounter_wave_direction'] = encounter_wave_direction_ds['encounter_wave_direction']
if mesh:
if bemio_import:
LOG.warning('Bemio data does not include mesh data. mesh=True is ignored.')
else:
body = results[0].body # Assumed to be all the same
dataset.coords['quadrature_method'] = str(body.mesh.quadrature_method)
dataset.coords["space_coordinate"] = xr.DataArray(["x", "y", "z"], dims=["space_coordinate"])
# `as_array_of_faces` is only defined on plain `Mesh`, not on a
# symmetric mesh (e.g. ReflectionSymmetricMesh/RotationSymmetricMesh)
# or a Multibody's joined mesh, so `.merged()` is used to get a plain
# mesh first. This does not change `nb_faces` or the order of the
# faces (matching `faces_centers` below), but it does mean any
# symmetry of the original mesh is not preserved in the dataset.
dataset["mesh_vertices"] = (["hull_face", "vertices_of_face", "space_coordinate"], body.mesh.merged().as_array_of_faces())
dataset["mesh_faces_center"] = (["hull_face", "space_coordinate"], body.mesh.faces_centers)
if body.lid_mesh is not None:
dataset["lid_mesh_vertices"] = (["lid_face", "vertices_of_face", "space_coordinate"], body.lid_mesh.merged().as_array_of_faces())
dataset["lid_mesh_faces_center"] = (["lid_face", "space_coordinate"], body.lid_mesh.faces_centers)
if len(body.dofs) > 0:
# `influenced_dof` (rather than `radiating_dof`) is used here because it
# always covers the full set of body.dofs and is present as soon as any
# radiation or diffraction result exists, whereas `radiating_dof` may be
# a user-restricted subset, or absent entirely for diffraction-only results.
if "influenced_dof" in dataset.coords:
dof_names = list(dataset.coords["influenced_dof"].values)
else:
dof_names = list(body.dofs.keys())
dof_motions = []
dof_gradient_of_motions = []
any_abstract_dof = False
for dof_name in dof_names:
dof = body.dofs[dof_name]
if isinstance(dof, AbstractDof):
any_abstract_dof = True
dof_motions.append(dof.evaluate_motion(body.mesh))
dof_gradient_of_motions.append(dof.evaluate_gradient_of_motion(body.mesh))
else:
dof_motions.append(dof)
dof_gradient_of_motions.append(np.full((body.mesh.nb_faces, 3, 3), np.nan))
mesh_dofs_dataset = xr.Dataset(
coords={"influenced_dof": dof_names},
data_vars={
"dof_motions": (["influenced_dof", "hull_face", "space_coordinate"], dof_motions),
},
)
if any_abstract_dof:
mesh_dofs_dataset.coords["gradient"] = xr.DataArray(["dx", "dy", "dz"], dims=["gradient"])
mesh_dofs_dataset["dof_gradient_of_motions"] = (
["influenced_dof", "hull_face", "space_coordinate", "gradient"], dof_gradient_of_motions)
dataset = xr.merge([dataset, mesh_dofs_dataset], compat="no_conflicts", join="outer")
# HYDROSTATICS
if hydrostatics:
if bemio_import:
LOG.warning('Bemio data import being used, hydrostatics=True is ignored.')
else:
body = results[0].body
if "radiating_dof" in dataset.coords:
radiating_dofs = dataset.coords["radiating_dof"].values.astype(str)
else:
radiating_dofs = None
try:
computed_hydrostatics = compute_hydrostatics_dataset(
body,
rho=dataset.coords["rho"].values,
g=dataset.coords["g"].values,
only_dofs=radiating_dofs,
)
except Exception as e:
LOG.warning(f"An error occurred while computing hydrostatics for body {body.__short_str__()}'.\n"
f"Error message: {e}\n"
f"Hydrostatics data for this body will be skipped. You can pass `hydrostatics=False` to `assemble_dataset` or `fill_dataset` to avoid this warning.")
computed_hydrostatics = xr.Dataset()
# If these dimensions are already in the dataset, we use the exact same coordinates to avoid issues when merging the datasets just below.
if "radiating_dof" in dataset.coords:
computed_hydrostatics = computed_hydrostatics.assign_coords(radiating_dof=dataset.coords["radiating_dof"].to_index())
if "influenced_dof" in dataset.coords:
computed_hydrostatics = computed_hydrostatics.assign_coords(influenced_dof=dataset.coords["influenced_dof"].to_index())
dataset = xr.merge([dataset, computed_hydrostatics], compat="no_conflicts", join="outer")
# EXPORT FIELDS ON BODY (IF KEPT IN THE RESULT OBJECT)
if not bemio_import:
# Only keep results with a free surface, consistently with `records` above.
results_with_free_surface = [r for r in results if r.free_surface == 0.0]
else:
results_with_free_surface = []
# No pressure on the lid: it is not a physical surface.
_add_face_distributions(dataset, results_with_free_surface, "pressure", optional_dims, lid=False)
if "pressure_diffraction" in dataset:
dataset["pressure_Froude_Krylov"] = _squeeze_dimensions(xr.merge(
[airy_waves_pressure_dataarray(r.problem) for r in results_with_free_surface if isinstance(r.problem, DiffractionProblem)],
compat="no_conflicts", join="outer"
)['pressure'], dimensions=optional_dims)
# Indirect method: sources. Direct method: potential (which is then the unknown of the linear system)
# (The potential is not stored for the indirect method, since it can be recomputed from the sources.)
_add_face_distributions(dataset, [r for r in results_with_free_surface if r.sources is not None],
"sources", optional_dims)
_add_face_distributions(dataset, [r for r in results_with_free_surface if r.sources is None],
"potential", optional_dims)
for var in set(dataset) | set(dataset.coords):
if var in VARIABLES_ATTRIBUTES:
dataset[var].attrs.update(VARIABLES_ATTRIBUTES[var])
dataset.attrs.update(attrs)
dataset.attrs['capytaine_version'] = __version__
return dataset
[docs]
def assemble_matrices(results):
"""Simplified version of assemble_dataset, returning only bare matrices.
Meant mainly for teaching without introducing Xarray to beginners.
Parameters
----------
results: list of LinearPotentialFlowResult
The results that will be read.
Returns
-------
3-ple of (np.arrays or None)
The added mass matrix, the radiation damping matrix and the excitation force.
If the data are no available in the results, returns None instead.
The matrices are ordered along their frequency axis in the same order as
the frequencies appear in ``results``, rather than sorted by increasing
frequency as in the dataset returned by :func:`assemble_dataset`.
"""
ds = assemble_dataset(results)
# `assemble_dataset` sorts the frequencies by increasing value. For the bare
# matrices, reorder the frequency axis to match the order in which the
# frequencies appear in `results`, which is less surprising when the matrices
# are read without the labelled coordinates of the dataset.
# See https://github.com/capytaine/capytaine/issues/797
freq_type = next(
(ft for ft in ('omega', 'freq', 'period', 'wavelength', 'wavenumber') if ft in ds.dims),
None,
)
if freq_type is not None:
ordered_freqs = list(dict.fromkeys(getattr(result, freq_type) for result in results))
ds = ds.sel({freq_type: ordered_freqs})
if "added_mass" in ds:
A = np.atleast_2d(ds.added_mass.values.squeeze())
else:
A = None
if "radiation_damping" in ds:
B = np.atleast_2d(ds.radiation_damping.values.squeeze())
else:
B = None
if "excitation_force" in ds:
F = np.atleast_1d(ds.excitation_force.values.squeeze())
else:
F = None
return A, B, F
################################
# Handling of complex values #
################################
[docs]
def separate_complex_values(ds: xr.Dataset) -> xr.Dataset:
"""Return a new Dataset where complex-valued arrays of shape (...)
have been replaced by real-valued arrays of shape (2, ...).
.. seealso::
:func:`merge_complex_values`
The invert operation
"""
ds = ds.copy()
for variable in ds.data_vars:
if ds[variable].dtype == complex:
da = ds[variable]
new_da = xr.DataArray(np.asarray((np.real(da).data, np.imag(da).data)),
dims=('complex',) + da.dims,
attrs=da.attrs)
ds[variable] = new_da
ds.coords['complex'] = ['re', 'im']
return ds
[docs]
def merge_complex_values(ds: xr.Dataset) -> xr.Dataset:
"""Return a new Dataset where real-valued arrays of shape (2, ...)
have been replaced by complex-valued arrays of shape (...).
.. seealso::
:func:`separate_complex_values`
The invert operation
"""
if 'complex' in ds.coords:
ds = ds.copy()
for variable in ds.data_vars:
if 'complex' in ds[variable].coords:
da = ds[variable]
new_dims = [d for d in da.dims if d != 'complex']
new_da = xr.DataArray(da.sel(complex='re').data + 1j*da.sel(complex='im').data, dims=new_dims)
ds[variable] = new_da
ds = ds.drop_vars('complex')
return ds
##################
# Save dataset #
##################
[docs]
def save_dataset_as_netcdf(filename, dataset):
"""Save `dataset` as a NetCDF file with name (or path) `filename`"""
ds = separate_complex_values(dataset)
# Workaround https://github.com/capytaine/capytaine/issues/683
ds['radiating_dof'] = ds['radiating_dof'].astype('str')
ds['influenced_dof'] = ds['influenced_dof'].astype('str')
# Make sure all strings are exported as strings and not Python objects
encoding = {'radiating_dof': {'dtype': 'U'},
'influenced_dof': {'dtype': 'U'}}
ds.to_netcdf(filename, encoding=encoding)
[docs]
def export_dataset(filename, dataset, format=None, **kwargs):
"""Save `dataset` into a format, provided by the `format` argument or inferred by the `filename`.
Parameters
----------
filename: str or Path
Where to store the data
dataset: xarray.Dataset
Dataset, which is assumed to have been computed by Capytaine
format: str, optional
Format of output. Accepted values: "netcdf"
**kwargs: optional
Remaining argument are passed to the specific export function,
such as ``save_dataset_as_netcdf``, ``export_to_wamit`` or ``write_dataset_as_tecplot_files``.
Returns
-------
None
"""
if (
(format is not None and format.lower() == "netcdf") or
(format is None and str(filename).endswith(".nc"))
):
save_dataset_as_netcdf(filename, dataset, **kwargs)
elif (
(format is not None and format.lower() == "wamit")
):
from capytaine.io.wamit import export_to_wamit
export_to_wamit(dataset, filename, **kwargs)
elif (
(format is not None and format.lower() == "nemoh")
):
from capytaine.io.legacy import write_dataset_as_tecplot_files, export_hydrostatics_from_dataset
Path(filename).mkdir(exist_ok=True)
try:
write_dataset_as_tecplot_files(filename, dataset, **kwargs)
except Exception as e:
LOG.warning(f"Export to Nemoh format: did not export hydrodynamics in {filename}: {e}")
try:
export_hydrostatics_from_dataset(filename, dataset)
except Exception as e:
LOG.warning(f"Export to Nemoh format: did not export hydrostatics in {filename}: {e}")
else:
raise ValueError("`export_dataset` could not infer export format based on filename or `format` argument.\n"
f"provided filename: {filename}\nprovided format: {format}")