Utilities

Module specs, logging and serialization

utils

Module construction, experiment logging, serialization and image operators.

JSONLogger

JSONLogger(path: str | PathLike[str] | None = None)

Logger that appends timestamped records to a JSON file.

Records are pretty-printed JSON objects separated by newlines. Read them with read_json_records or iter_json_records.

Parameters:

• path (str | PathLike[str] | None, default: None ) –

  File to append to; None or an empty path disables writing.

Source code in inro/utils.py

def __init__(self, path: str | os.PathLike[str] | None = None) -> None:
    self.path = path

log

log(entry: Mapping[str, Any]) -> None

Append a JSON-serializable record with a UTC ISO 8601 timestamp.

Parameters:

• entry (Mapping[str, Any]) –

  Record to append; its own timestamp, if any, takes precedence.

Source code in inro/utils.py

def log(self, entry: Mapping[str, Any]) -> None:
    """Append a JSON-serializable record with a UTC ISO 8601 timestamp.

    Args:
        entry: Record to append; its own ``timestamp``, if any, takes precedence.
    """
    timestamp = datetime.now(timezone.utc).isoformat()
    entry = {"timestamp": timestamp, **entry}
    if self.path:
        with open(self.path, "a", encoding="utf-8") as log_file:
            log_file.write(json.dumps(entry, ensure_ascii=False, indent=4) + "\n")

build_module

build_module(spec: ModuleSpec, registry: Mapping[str, type[Module]], *args: Any, **kwargs: Any) -> Module

Build a module from an instance, registry key, class, or (key/class, params).

An instance is returned unchanged. Otherwise the class receives *args, **params and **kwargs.

Parameters:

• spec (ModuleSpec) –

  Module instance, key, class, or a key or class paired with parameters, in a tuple or a list, so that specs can come from JSON.

• registry (Mapping[str, type[Module]]) –

  Mapping from string keys to module classes.

• *args (Any, default: () ) –

  Positional arguments of the constructor.

• **kwargs (Any, default: {} ) –

  Keyword arguments of the constructor, which must not repeat params.

Raises:

• ValueError –

  If spec has another form.

• KeyError –

  If the key is not in registry.

Examples:

>>> build_module(("linear", {"bias": False}), {"linear": nn.Linear}, 4, 2)
Linear(in_features=4, out_features=2, bias=False)

Source code in inro/utils.py

def build_module(
    spec: ModuleSpec,
    registry: Mapping[str, type[nn.Module]],
    *args: Any,
    **kwargs: Any,
) -> nn.Module:
    """Build a module from an instance, registry key, class, or ``(key/class, params)``.

    An instance is returned unchanged. Otherwise the class receives ``*args``,
    ``**params`` and ``**kwargs``.

    Args:
        spec: Module instance, key, class, or a key or class paired with parameters, in
            a tuple or a list, so that specs can come from JSON.
        registry: Mapping from string keys to module classes.
        *args: Positional arguments of the constructor.
        **kwargs: Keyword arguments of the constructor, which must not repeat
            ``params``.

    Raises:
        ValueError: If ``spec`` has another form.
        KeyError: If the key is not in ``registry``.

    Examples:
        >>> build_module(("linear", {"bias": False}), {"linear": nn.Linear}, 4, 2)
        Linear(in_features=4, out_features=2, bias=False)
    """
    if isinstance(spec, nn.Module):
        return spec

    if isinstance(spec, (str, type)):
        key_or_class, params = spec, {}
    elif isinstance(spec, (tuple, list)) and len(spec) == 2:
        key_or_class, params = spec
    else:
        raise ValueError(
            "spec must be an nn.Module instance, a registry key / nn.Module "
            f"subclass, or a (name/type, params) pair, got {spec!r}"
        )

    if isinstance(key_or_class, str):
        if key_or_class not in registry:
            known = ", ".join(sorted(registry)) or "none"
            raise KeyError(
                f"unknown key {key_or_class!r}; known keys: {known}. Pass an "
                "nn.Module subclass or instance instead."
            )
        module_class = registry[key_or_class]
    else:
        module_class = key_or_class
    return module_class(*args, **params, **kwargs)

count_parameters

count_parameters(model: Module, trainable_only: bool = True) -> int

Count parameters as real numbers, counting complex elements twice.

Counting real components makes real and complex models comparable.

Parameters:

• model (Module) –

  Module whose parameters are counted.

• trainable_only (bool, default: True ) –

  Whether to count only parameters with requires_grad=True.

Examples:

>>> count_parameters(nn.Linear(2, 3))
9

Source code in inro/utils.py

def count_parameters(model: nn.Module, trainable_only: bool = True) -> int:
    """Count parameters as real numbers, counting complex elements twice.

    Counting real components makes real and complex models comparable.

    Args:
        model: Module whose parameters are counted.
        trainable_only: Whether to count only parameters with ``requires_grad=True``.

    Examples:
        >>> count_parameters(nn.Linear(2, 3))
        9
    """
    return sum(
        param.numel() * (2 if param.is_complex() else 1)
        for param in model.parameters()
        if param.requires_grad or not trainable_only
    )

setup_experiment_loggers

setup_experiment_loggers(log_dir: str | PathLike[str] | None, logger_name: str, log_filename: str = 'log.txt', json_filename: str = 'log.json', mode: Literal['w', 'a'] = 'w') -> tuple[Logger, JSONLogger]

Configure console, text and JSON logging for an experiment.

Replaces the named logger's handlers and disables propagation. The console shows INFO messages, written around a progress bar rather than through it, and the text file DEBUG ones as well. JSON records always append.

Parameters:

• log_dir (str | PathLike[str] | None) –

  Directory of the logs, created if needed; None writes no files.

• logger_name (str) –

  Name passed to logging.getLogger.

• log_filename (str, default: 'log.txt' ) –

  Name of the text log within log_dir.

• json_filename (str, default: 'log.json' ) –

  Name of the JSON log within log_dir.

• mode (Literal['w', 'a'], default: 'w' ) –

  "w" truncates the text log first; "a" appends, so that training and evaluation can share the file.

Returns:

• tuple[Logger, JSONLogger] –

  The text logger and the JSON logger, which writes nothing without log_dir.

Source code in inro/utils.py

def setup_experiment_loggers(
    log_dir: str | os.PathLike[str] | None,
    logger_name: str,
    log_filename: str = "log.txt",
    json_filename: str = "log.json",
    mode: Literal["w", "a"] = "w",
) -> tuple[logging.Logger, JSONLogger]:
    """Configure console, text and JSON logging for an experiment.

    Replaces the named logger's handlers and disables propagation. The console shows
    ``INFO`` messages, written around a progress bar rather than through it, and the
    text file ``DEBUG`` ones as well. JSON records always append.

    Args:
        log_dir: Directory of the logs, created if needed; ``None`` writes no files.
        logger_name: Name passed to :func:`logging.getLogger`.
        log_filename: Name of the text log within ``log_dir``.
        json_filename: Name of the JSON log within ``log_dir``.
        mode: ``"w"`` truncates the text log first; ``"a"`` appends, so that training
            and evaluation can share the file.

    Returns:
        The text logger and the JSON logger, which writes nothing without ``log_dir``.
    """
    if log_dir:
        os.makedirs(log_dir, exist_ok=True)
        log_path = os.path.join(log_dir, log_filename)
        json_log_path = os.path.join(log_dir, json_filename)
    else:
        log_path = None
        json_log_path = None

    logger = logging.getLogger(logger_name)
    logger.handlers.clear()
    logger.setLevel(logging.DEBUG)
    logger.propagate = False

    formatter = logging.Formatter(
        fmt="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
    )

    stream_handler = _ProgressBarHandler()
    stream_handler.setLevel(logging.INFO)
    stream_handler.setFormatter(formatter)
    logger.addHandler(stream_handler)

    if log_path:
        if mode == "w":
            # Append thereafter so training and evaluation can share the file.
            open(log_path, "w", encoding="utf-8").close()
        file_handler = logging.FileHandler(log_path, mode="a", encoding="utf-8")
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)

    json_logger = JSONLogger(json_log_path)
    return logger, json_logger

iter_json_records

iter_json_records(path: str | PathLike[str]) -> Iterator[dict[str, Any]]

Yield records from a stream of JSON objects, including pretty-printed logs.

Parsing stops at the first malformed or incomplete record, so an interrupted final write is ignored.

Parameters:

• path (str | PathLike[str]) –

  JSON log file.

Yields:

• dict[str, Any] –

  Each record as a dictionary.

Source code in inro/utils.py

def iter_json_records(path: str | os.PathLike[str]) -> Iterator[dict[str, Any]]:
    """Yield records from a stream of JSON objects, including pretty-printed logs.

    Parsing stops at the first malformed or incomplete record, so an interrupted final
    write is ignored.

    Args:
        path: JSON log file.

    Yields:
        Each record as a dictionary.
    """
    decoder = json.JSONDecoder()
    with open(path, "r", encoding="utf-8") as log_file:
        buffer = ""
        for line in log_file:
            buffer += line
            while buffer:
                buffer = buffer.lstrip()
                try:
                    record, end = decoder.raw_decode(buffer)
                    yield record
                    buffer = buffer[end:]
                except json.JSONDecodeError:
                    break

read_json_records

read_json_records(path: str | PathLike[str]) -> list[dict[str, Any]]

Read all records that iter_json_records yields.

Parameters:

• path (str | PathLike[str]) –

  JSON log file.

Source code in inro/utils.py

def read_json_records(path: str | os.PathLike[str]) -> list[dict[str, Any]]:
    """Read all records that :func:`iter_json_records` yields.

    Args:
        path: JSON log file.
    """
    return list(iter_json_records(path))

sobel_gradients

sobel_gradients(image: Tensor, normalize: bool = True, padding_mode: Literal['constant', 'reflect', 'replicate', 'circular'] = 'replicate') -> tuple[Tensor, Tensor]

Return per-channel Sobel gradients (gx, gy) of a channels-last batch.

gx points along increasing columns and gy along increasing rows. Computation uses float32.

Parameters:

• image (Tensor) –

  Image batch.

• normalize (bool, default: True ) –

  Whether to divide the kernels by 8, so that a unit ramp has unit gradient.

• padding_mode (Literal['constant', 'reflect', 'replicate', 'circular'], default: 'replicate' ) –

  Mode of the one-pixel padding; "constant" pads with zeros.

Shape

• Input: (B, H, W, channels).
• Output: Two tensors (B, H, W, channels).

Examples:

>>> ramp = torch.arange(4.0).expand(4, 4)[None, :, :, None]
>>> gx, gy = sobel_gradients(ramp)
>>> gx[0, 1, 1, 0].item(), gy[0, 1, 1, 0].item()
(1.0, 0.0)

Source code in inro/utils.py

def sobel_gradients(
    image: Tensor,
    normalize: bool = True,
    padding_mode: Literal["constant", "reflect", "replicate", "circular"] = "replicate",
) -> tuple[Tensor, Tensor]:
    """Return per-channel Sobel gradients ``(gx, gy)`` of a channels-last batch.

    ``gx`` points along increasing columns and ``gy`` along increasing rows. Computation
    uses float32.

    Args:
        image: Image batch.
        normalize: Whether to divide the kernels by 8, so that a unit ramp has unit
            gradient.
        padding_mode: Mode of the one-pixel padding; ``"constant"`` pads with zeros.

    Shape:
        - Input: ``(B, H, W, channels)``.
        - Output: Two tensors ``(B, H, W, channels)``.

    Examples:
        >>> ramp = torch.arange(4.0).expand(4, 4)[None, :, :, None]
        >>> gx, gy = sobel_gradients(ramp)
        >>> gx[0, 1, 1, 0].item(), gy[0, 1, 1, 0].item()
        (1.0, 0.0)
    """
    x = image.movedim(-1, 1).float()
    channels = x.shape[1]

    gx_kernel = torch.tensor(
        [[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]],
        dtype=x.dtype,
        device=x.device,
    )
    if normalize:
        gx_kernel = gx_kernel / 8.0

    padded = F.pad(x, (1, 1, 1, 1), mode=padding_mode)
    gradients = [
        F.conv2d(padded, kernel.expand(channels, 1, 3, 3), groups=channels)
        for kernel in (gx_kernel, gx_kernel.t())
    ]
    return gradients[0].movedim(1, -1), gradients[1].movedim(1, -1)

encode_dict

encode_dict(data: dict[str, Any], compress: bool = True) -> bytes

Serialize a tensor dictionary to bytes with torch.save.

To round-trip through decode_dict, values must be tensors, primitive types or standard containers accepted by torch.load(weights_only=True).

Parameters:

• data (dict[str, Any]) –

  Dictionary to serialize.

• compress (bool, default: True ) –

  Whether to compress with gzip.

Examples:

>>> encoded = encode_dict({"weight": torch.ones(2)})
>>> decode_dict(encoded)["weight"].tolist()
[1.0, 1.0]

Source code in inro/utils.py

def encode_dict(data: dict[str, Any], compress: bool = True) -> bytes:
    """Serialize a tensor dictionary to bytes with :func:`torch.save`.

    To round-trip through :func:`decode_dict`, values must be tensors, primitive types
    or standard containers accepted by ``torch.load(weights_only=True)``.

    Args:
        data: Dictionary to serialize.
        compress: Whether to compress with gzip.

    Examples:
        >>> encoded = encode_dict({"weight": torch.ones(2)})
        >>> decode_dict(encoded)["weight"].tolist()
        [1.0, 1.0]
    """
    buffer = io.BytesIO()
    torch.save(data, buffer)
    raw = buffer.getvalue()
    return gzip.compress(raw) if compress else raw

decode_dict

decode_dict(encoded: bytes, compressed: bool = True) -> dict[str, Any]

Deserialize encode_dict output, loading tensors onto the CPU.

Loading uses PyTorch's default weights_only=True.

Parameters:

• encoded (bytes) –

  Serialized dictionary.

• compressed (bool, default: True ) –

  Whether encoded is gzipped, matching the encoder's compress.

Source code in inro/utils.py

def decode_dict(encoded: bytes, compressed: bool = True) -> dict[str, Any]:
    """Deserialize :func:`encode_dict` output, loading tensors onto the CPU.

    Loading uses PyTorch's default ``weights_only=True``.

    Args:
        encoded: Serialized dictionary.
        compressed: Whether ``encoded`` is gzipped, matching the encoder's ``compress``.
    """
    raw = gzip.decompress(encoded) if compressed else encoded
    return torch.load(io.BytesIO(raw), map_location="cpu")

read_config

read_config(path: str | PathLike[str]) -> dict[str, Any]

Read a UTF-8 JSON configuration file.

Parameters:

• path (str | PathLike[str]) –

  File to read.

Source code in inro/utils.py

def read_config(path: str | os.PathLike[str]) -> dict[str, Any]:
    """Read a UTF-8 JSON configuration file.

    Args:
        path: File to read.
    """
    with open(path, "r", encoding="utf-8") as json_file:
        return json.load(json_file)

save_results

save_results(results: Sequence[dict[str, Any]], save_dir: str | PathLike[str], prefix: str = '') -> None

Write records to save_dir/{prefix}_results.json, replacing any existing file.

Parameters:

• results (Sequence[dict[str, Any]]) –

  JSON-serializable records.

• save_dir (str | PathLike[str]) –

  Output directory, created if needed.

• prefix (str, default: '' ) –

  Prefix of the file name.

Source code in inro/utils.py

def save_results(
    results: Sequence[dict[str, Any]],
    save_dir: str | os.PathLike[str],
    prefix: str = "",
) -> None:
    """Write records to ``save_dir/{prefix}_results.json``, replacing any existing file.

    Args:
        results: JSON-serializable records.
        save_dir: Output directory, created if needed.
        prefix: Prefix of the file name.
    """
    os.makedirs(save_dir, exist_ok=True)
    with open(
        os.path.join(save_dir, f"{prefix}_results.json"), "w", encoding="utf-8"
    ) as json_file:
        json.dump(results, json_file, indent=4)

add_sensor_noise

add_sensor_noise(image: Tensor, readout_noise_std: float = 1.0, integration_time: float = 1.0, seed: int = 0) -> Tensor

Simulate photon shot noise and Gaussian readout noise.

The image is scaled to photon counts, Poisson-sampled and perturbed by Gaussian noise, then divided by the integration time and clamped to [0, 1].

Parameters:

• image (Tensor) –

  Image in [0, 1].

• readout_noise_std (float, default: 1.0 ) –

  Gaussian standard deviation in count units.

• integration_time (float, default: 1.0 ) –

  Photon-count scale; longer exposures reduce relative noise.

• seed (int, default: 0 ) –

  Seed of the global PyTorch RNG.

Returns:

• Tensor –

  Noisy float32 image, with the shape and device of image.

Source code in inro/utils.py

def add_sensor_noise(
    image: Tensor,
    readout_noise_std: float = 1.0,
    integration_time: float = 1.0,
    seed: int = 0,
) -> Tensor:
    """Simulate photon shot noise and Gaussian readout noise.

    The image is scaled to photon counts, Poisson-sampled and perturbed by Gaussian
    noise, then divided by the integration time and clamped to ``[0, 1]``.

    Args:
        image: Image in ``[0, 1]``.
        readout_noise_std: Gaussian standard deviation in count units.
        integration_time: Photon-count scale; longer exposures reduce relative noise.
        seed: Seed of the global PyTorch RNG.

    Returns:
        Noisy float32 image, with the shape and device of ``image``.
    """
    torch.manual_seed(seed)
    # Sample in float32 even when the input uses a lower-precision dtype.
    shot_noise = torch.poisson(image.float() * integration_time)
    readout_noise = torch.randn_like(shot_noise) * readout_noise_std
    noisy = (shot_noise + readout_noise) / integration_time
    return noisy.clamp(0.0, 1.0)

render_mesh

render_mesh(mesh_path: str | PathLike[str], image_path: str | PathLike[str], width: int = 1920, height: int = 1080) -> None

Render a mesh to an image with a hidden Open3D window (3d extra).

Parameters:

• mesh_path (str | PathLike[str]) –

  Mesh file loaded with trimesh; a scene is joined into one mesh.

• image_path (str | PathLike[str]) –

  Output image, whose suffix sets the format.

• width (int, default: 1920 ) –

  Image width in pixels.

• height (int, default: 1080 ) –

  Image height in pixels.

Source code in inro/utils.py

def render_mesh(
    mesh_path: str | os.PathLike[str],
    image_path: str | os.PathLike[str],
    width: int = 1920,
    height: int = 1080,
) -> None:
    """Render a mesh to an image with a hidden Open3D window (``3d`` extra).

    Args:
        mesh_path: Mesh file loaded with trimesh; a scene is joined into one mesh.
        image_path: Output image, whose suffix sets the format.
        width: Image width in pixels.
        height: Image height in pixels.
    """
    import open3d as o3d
    import trimesh

    mesh = trimesh.load(mesh_path)
    if isinstance(mesh, trimesh.Scene):
        mesh = mesh.dump(concatenate=True)

    o3d_mesh = o3d.geometry.TriangleMesh()
    o3d_mesh.vertices = o3d.utility.Vector3dVector(mesh.vertices)
    o3d_mesh.triangles = o3d.utility.Vector3iVector(mesh.faces)
    o3d_mesh.compute_vertex_normals()

    visualizer = o3d.visualization.Visualizer()
    visualizer.create_window(width=width, height=height, visible=False)
    visualizer.add_geometry(o3d_mesh)
    visualizer.capture_screen_image(str(image_path), do_render=True)
    visualizer.destroy_window()

radon

radon(image: Tensor, angles: Tensor) -> Tensor

Compute a sinogram by rotating the image and summing along its height.

Rotations use Kornia's bilinear interpolation and zero padding.

Parameters:

• image (Tensor) –

  Image (1, channels, H, W).

• angles (Tensor) –

  Angles (A,) in degrees, on the image's device.

Returns:

• Tensor –

  Sinogram (A, channels, W), with singleton dimensions removed.

Source code in inro/utils.py

def radon(image: Tensor, angles: Tensor) -> Tensor:
    """Compute a sinogram by rotating the image and summing along its height.

    Rotations use Kornia's bilinear interpolation and zero padding.

    Args:
        image: Image ``(1, channels, H, W)``.
        angles: Angles ``(A,)`` in degrees, on the image's device.

    Returns:
        Sinogram ``(A, channels, W)``, with singleton dimensions removed.
    """
    import kornia

    rotated = kornia.geometry.rotate(image.repeat_interleave(len(angles), 0), angles)
    return rotated.sum(2).squeeze()

Plotting

visualization

Plotting helpers for experiment results.

LinePlot

LinePlot(data: DataFrame | Mapping[str, Sequence[Any]] | Sequence[dict[str, Any]], columns: Sequence[str] | None = None)

Seaborn line plot of tabular results, such as rate-distortion curves.

Parameters:

• data (DataFrame | Mapping[str, Sequence[Any]] | Sequence[dict[str, Any]]) –

  Results in any form pandas.DataFrame accepts: a DataFrame, a mapping from column names to sequences, or a sequence of record dictionaries.

• columns (Sequence[str] | None, default: None ) –

  Columns passed to pandas.DataFrame; None keeps all columns.

Attributes:

• data –

  The results as a pandas.DataFrame.

• fig (Figure | None) –

  Figure created by the last plot call, or None.

• ax (Axes | None) –

  Axes created by the last plot call, or None.

Examples:

>>> records = [
...     {"bpp": 0.1, "psnr": 30.0, "method": "FUTON"},
...     {"bpp": 0.2, "psnr": 33.0, "method": "FUTON"},
... ]
>>> plot = LinePlot(records)
>>> fig, ax = plot.plot(x="bpp", y="psnr", hue="method")
>>> plot.save("figures", "rate_distortion")

Source code in inro/visualization.py

def __init__(
    self,
    data: pd.DataFrame | Mapping[str, Sequence[Any]] | Sequence[dict[str, Any]],
    columns: Sequence[str] | None = None,
) -> None:
    self.data = pd.DataFrame(data, columns=columns)
    self.fig: Figure | None = None
    self.ax: Axes | None = None

plot

plot(x: str, y: str, hue: str, xlim: tuple[float | None, float | None] = (None, None), ylim: tuple[float | None, float | None] = (None, None), legend_labels: Sequence[str] | None = None, legend_loc: str = 'lower right', **kwargs: Any) -> tuple[Figure, Axes]

Draw lines grouped by hue and store the resulting fig and ax.

Parameters:

• x (str) –

  Column of the x-axis.

• y (str) –

  Column of the y-axis.

• hue (str) –

  Column defining both line color and style.

• xlim (tuple[float | None, float | None], default: (None, None) ) –

  Left and right limits; None keeps a limit automatic.

• ylim (tuple[float | None, float | None], default: (None, None) ) –

  Bottom and top limits; None keeps a limit automatic.

• legend_labels (Sequence[str] | None, default: None ) –

  Existing labels to keep in the legend, in order, matching exactly; None uses the hue values in order of appearance, which must be strings.

• legend_loc (str, default: 'lower right' ) –

  Legend location accepted by matplotlib.axes.Axes.legend.

• **kwargs (Any, default: {} ) –

  Passed to seaborn.lineplot, such as aggregation options.

Returns:

• tuple[Figure, Axes] –

  The new figure and axes.

Source code in inro/visualization.py

def plot(
    self,
    x: str,
    y: str,
    hue: str,
    xlim: tuple[float | None, float | None] = (None, None),
    ylim: tuple[float | None, float | None] = (None, None),
    legend_labels: Sequence[str] | None = None,
    legend_loc: str = "lower right",
    **kwargs: Any,
) -> tuple[Figure, Axes]:
    """Draw lines grouped by ``hue`` and store the resulting ``fig`` and ``ax``.

    Args:
        x: Column of the x-axis.
        y: Column of the y-axis.
        hue: Column defining both line color and style.
        xlim: Left and right limits; ``None`` keeps a limit automatic.
        ylim: Bottom and top limits; ``None`` keeps a limit automatic.
        legend_labels: Existing labels to keep in the legend, in order, matching
            exactly; ``None`` uses the ``hue`` values in order of appearance, which
            must be strings.
        legend_loc: Legend location accepted by ``matplotlib.axes.Axes.legend``.
        **kwargs: Passed to ``seaborn.lineplot``, such as aggregation options.

    Returns:
        The new figure and axes.
    """
    if legend_labels is None:
        legend_labels = self.data[hue].unique()

    fig, ax = plt.subplots()
    sns.lineplot(
        ax=ax,
        data=self.data,
        x=x,
        y=y,
        hue=hue,
        style=hue,
        legend="brief",
        **kwargs,
    )
    ax.grid()
    ax.set_xlim(*xlim)
    ax.set_ylim(*ylim)

    handles, labels = ax.get_legend_handles_labels()
    selected_handles = [handles[labels.index(label)] for label in legend_labels]
    ax.legend(selected_handles, list(legend_labels), loc=legend_loc, title="")

    self.fig, self.ax = fig, ax
    return fig, ax

save

save(save_dir: str | PathLike[str] = '.', file_name: str = '', format: str = 'pdf') -> None

Save the last plot to save_dir/{file_name}.{format} with tight bounds.

Call plot first.

Parameters:

• save_dir (str | PathLike[str], default: '.' ) –

  Output directory, created if needed.

• file_name (str, default: '' ) –

  File name without extension; must not be empty, for matplotlib to recognize the extension.

• format (str, default: 'pdf' ) –

  Extension of the file, which sets its format.

Source code in inro/visualization.py

def save(
    self,
    save_dir: str | os.PathLike[str] = ".",
    file_name: str = "",
    format: str = "pdf",
) -> None:
    """Save the last plot to ``save_dir/{file_name}.{format}`` with tight bounds.

    Call :meth:`plot` first.

    Args:
        save_dir: Output directory, created if needed.
        file_name: File name without extension; must not be empty, for matplotlib to
            recognize the extension.
        format: Extension of the file, which sets its format.
    """
    os.makedirs(save_dir, exist_ok=True)

    self.fig.savefig(
        os.path.join(save_dir, f"{file_name}.{format}"),
        bbox_inches="tight",
        pad_inches=0,
    )