Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Some plotting improvements #717

Merged
merged 9 commits into from
Apr 8, 2024
32 changes: 20 additions & 12 deletions pastas/plotting/modelcompare.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

from itertools import combinations
from logging import getLogger
from typing import List, Optional, Tuple
from typing import List, Optional
from warnings import warn

import matplotlib.pyplot as plt
Expand Down Expand Up @@ -92,9 +92,9 @@ def __init__(self, models: List[Model], names: Optional[List[str]] = None) -> No
def initialize_figure(
self,
mosaic: Optional[List[List[str]]] = None,
figsize: Tuple[int, int] = (10, 8),
cmap: str = "tab10",
return_ax: bool = False,
**fig_kwargs,
) -> None:
"""initialize a custom figure based on a mosaic.

Expand All @@ -112,7 +112,7 @@ def initialize_figure(

self.cmap = plt.get_cmap(cmap)

figure, axes = plt.subplot_mosaic(mosaic, figsize=figsize)
figure, axes = plt.subplot_mosaic(mosaic, **fig_kwargs)
if return_ax:
return axes

Expand All @@ -123,9 +123,9 @@ def initialize_figure(
def initialize_adjust_height_figure(
self,
mosaic: Optional[List[List[str]]] = None,
figsize: Tuple[int] = (10, 8),
cmap: str = "tab10",
smdict: Optional[dict] = None,
**fig_kwargs,
) -> None:
"""initialize subplots based on a mosaic with equal vertical scales.

Expand Down Expand Up @@ -226,8 +226,8 @@ def initialize_adjust_height_figure(
self.mosaic = mosaic
fig, axes = plt.subplot_mosaic(
self.mosaic,
figsize=figsize,
gridspec_kw=dict(height_ratios=heights_list),
**fig_kwargs,
)

self.figure = fig
Expand Down Expand Up @@ -577,10 +577,14 @@ def plot_response(
continue
if response == "step":
kwargs = {}
p = None
if ml.stressmodels[smn].rfunc is not None:
if isinstance(ml.stressmodels[smn].rfunc, HantushWellModel):
kwargs = {"warn": False}
step = ml.get_step_response(smn, add_0=True, **kwargs)
p = ml.stressmodels[smn].get_parameters(
model=ml, istress=0
)
step = ml.get_step_response(smn, p=p, add_0=True, **kwargs)
if step is None:
continue
if self.axes is None:
Expand All @@ -598,10 +602,14 @@ def plot_response(
)
elif response == "block":
kwargs = {}
p = None
if ml.stressmodels[smn].rfunc is not None:
if isinstance(ml.stressmodels[smn].rfunc, HantushWellModel):
kwargs = {"warn": False}
block = ml.get_block_response(smn, **kwargs)
p = ml.stressmodels[smn].get_parameters(
model=ml, istress=0
)
block = ml.get_block_response(smn, p=p, add_0=True, **kwargs)
if block is None:
continue
if self.axes is None:
Expand Down Expand Up @@ -897,11 +905,11 @@ def plot(
smdict: Optional[dict] = None,
normalized: bool = False,
param_selection: Optional[list] = None,
figsize: Optional[tuple] = (10, 8),
grid: bool = True,
legend: bool = True,
adjust_height: bool = False,
legend_kwargs: Optional[dict] = None,
**fig_kwargs,
) -> None:
"""plot the models in a comparison plot.

Expand All @@ -921,8 +929,6 @@ def plot(
zero, by default False.
param_selection : list, optional
list of (sub)strings of which parameters to show in table, by default None.
figsize : tuple, optional
figure size, by default (10, 8).
grid : bool, optional
grid in each subplot, by default True.
legend : bool, optional
Expand All @@ -935,10 +941,12 @@ def plot(
pass legend keyword arguments to plots.
"""
self.adjust_height = adjust_height
if "figsize" not in fig_kwargs:
fig_kwargs["figsize"] = (10, 8)
if self.axes is None and not self.adjust_height:
self.initialize_figure(figsize=figsize)
self.initialize_figure(**fig_kwargs)
if self.axes is None and self.adjust_height:
self.initialize_adjust_height_figure(smdict=smdict, figsize=figsize)
self.initialize_adjust_height_figure(smdict=smdict, **fig_kwargs)

# sim
_ = self.plot_oseries()
Expand Down
10 changes: 7 additions & 3 deletions pastas/plotting/modelplots.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
"""This module contains all the plotting methods for Pastas Models."""
"""This module contains plotting methods for Pastas Models."""

import logging

Expand Down Expand Up @@ -267,11 +267,15 @@ def results(

# plot the step response
rkwargs = {}
p = None
if self.ml.stressmodels[sm_name].rfunc is not None:
if isinstance(self.ml.stressmodels[sm_name].rfunc, HantushWellModel):
rkwargs = {"warn": False}
p = self.ml.stressmodels[sm_name].get_parameters(
model=self.ml, istress=0
)
response = self.ml._get_response(
block_or_step=block_or_step, name=sm_name, add_0=True, **rkwargs
block_or_step=block_or_step, name=sm_name, p=p, add_0=True, **rkwargs
)

if response is not None:
Expand Down Expand Up @@ -907,7 +911,7 @@ def custom_sort(t):
contributions.append((name, h))

# plot step responses for each well, scaled with distance
p = sml.get_parameters(istress=istress)
p = sml.get_parameters(model=self.ml, istress=istress)
step = self.ml.get_step_response(sm, p=p)
ax_step.plot(step.index, step, c=stackcolors[name], label=name)
# recalculate y-limits step response axes
Expand Down
14 changes: 13 additions & 1 deletion pastas/plotting/plotly.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
"""This module contains interactive plots for Pastas models.
"""

import numpy as np
import pandas as pd
import plotly.graph_objs as go
Expand All @@ -10,6 +13,7 @@
_table_formatter_params,
_table_formatter_stderr,
)
from pastas.rfunc import HantushWellModel
from pastas.stats import acf


Expand Down Expand Up @@ -246,8 +250,16 @@ def results(self, tmin=None, tmax=None):
traces.append(trace_c)

# response
rkwargs = {}
p = None
if self._model.stressmodels[c.name].rfunc is not None:
if isinstance(self._model.stressmodels[c.name].rfunc, HantushWellModel):
rkwargs = {"warn": False}
p = self._model.stressmodels[c.name].get_parameters(
model=self._model, istress=0
)
response = self._model._get_response(
block_or_step="step", name=c.name, add_0=True
block_or_step="step", name=c.name, p=p, add_0=True, **rkwargs
)
trace_r = go.Scatter(
x=response.index,
Expand Down
2 changes: 1 addition & 1 deletion pastas/plotting/plots.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
"""This module contains all the plotting methods in Pastas."""
"""This module contains plotting methods for Pastas."""

import logging
from typing import Dict, List, Optional, Tuple, Union
Expand Down
3 changes: 3 additions & 0 deletions pastas/plotting/plotutil.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
"""This module contains utility functions for plotting.
"""

from typing import List, Union

import numpy as np
Expand Down
4 changes: 2 additions & 2 deletions pastas/timeseries_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@ def timestep_weighted_resample(s: Series, index: Index, fast: bool = False) -> S
# calculate the cumulative sum
s_new = s_new.cumsum()

# add NaNs at none-existing values in series at index
# add NaNs at non-existing values in series at index
s_new = s_new.combine_first(Series(np.NaN, index))

# interpolate these NaN's, only keep values at index
Expand All @@ -276,7 +276,7 @@ def timestep_weighted_resample(s: Series, index: Index, fast: bool = False) -> S
# calculate the diff again (inverse of cumsum)
s_new = s_new.diff()

# devide by the timestep again
# divide by the timestep again
s_new = s_new / _get_dt_array(s_new.index)

# set values after the end of the original series to NaN
Expand Down