hyrax.verbs

Submodules

• hyrax.verbs.database_connection
• hyrax.verbs.engine
• hyrax.verbs.infer
• hyrax.verbs.lookup
• hyrax.verbs.model
• hyrax.verbs.prepare
• hyrax.verbs.save_to_database
• hyrax.verbs.test
• hyrax.verbs.to_onnx
• hyrax.verbs.train
• hyrax.verbs.umap
• hyrax.verbs.verb_registry
• hyrax.verbs.visualize
• hyrax.verbs.visualize_v2

Classes

DatabaseConnection	Verb to create a connection to a vector database with inference results.
Umap	Umap latent space points into 2d
Infer	Inference verb
Train	Train verb
Test	Test verb - evaluates a trained model on test data
Visualize	Verb to create a visualization
VisualizeV2	Verb to create a hexbin visualization of a 2D latent space.
Lookup	Look up an inference result using the ID of a data member
SaveToDatabase	Verb to insert inference results into a vector database index for fast
Model	Resolves the model class that is defined in the config file.
ToOnnx	Export the model to ONNX format
Engine	This verb drives inference with an ONNX model in production.
Prepare	Prepare Verb, Prepares a dataset and returns it
Verb	Base class for all hyrax verbs

Functions

all_class_verbs(→ list[str])	Returns all verbs that are currently registered with a class-based implementation
all_verbs(→ list[str])	Returns all verbs that are currently registered
fetch_verb_class(→ type[Verb] | None)	Gives the class object for the named verb
is_verb_class(→ bool)	Returns true if the verb has a class based implementation

Package Contents

class DatabaseConnection(config)

Bases: hyrax.verbs.verb_registry.Verb

Verb to create a connection to a vector database with inference results.

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'database_connection'

add_parser_kwargs

description = 'Create a connection to the vector database for interactive queries.'

static setup_parser(parser: argparse.ArgumentParser)

Stub of parser setup

run_cli(args: argparse.Namespace | None = None)

Stub CLI implementation

run(database_dir: pathlib.Path | str | None = None)

Create a connection to the vector database for interactive queries.

Parameters:

database_dir (str or Path, Optional) – The directory containing the database that will be connected to. If None, attempt to connect to the most recently created …-vector-db-… directory. If specified, it can point to either an empty directory or a directory containing an existing vector database. If the latter, the database will be updated with the new vectors.

_get_database_type_from_config(database_dir: pathlib.Path)

Internal function that will read a config file from a directory and return the name of the vector database from it. i.e. “chromadb”, “qdrant”.

Parameters:

database_dir (Path) – The directory containing the vector database and the config file that be used as reference.

Returns:

The config value for [“vector_db”][“name”] in the reference config.

Return type:

str

class Umap(config)

Bases: hyrax.verbs.verb_registry.Verb

Umap latent space points into 2d

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'umap'

add_parser_kwargs

description = 'Transforms the entire dataset into a lower-dimensional space by fitting a UMAP model.'

static setup_parser(parser: argparse.ArgumentParser)

Stub of parser setup

run_cli(args: argparse.Namespace | None = None)

Stub CLI implementation

run(input_dir: pathlib.Path | str | None = None, model_path: pathlib.Path | str | None = None)

Create a umap of a particular inference run

This method loads the latent space representations from an inference run, samples a subset of data points, flattens them if necessary, and then fits a UMAP model. The fitted reducer is then used to transform the entire dataset into a lower-dimensional space.

Parameters:

• input_dir (str or Path, Optional) – The directory containing the inference results.

• model_path (str or Path, Optional) – The path to a pre-existing UMAP model. If provided, the method will use this model instead of fitting a new one.

Returns:

The method does not return anything but saves the UMAP representations to disk.

Return type:

None

_run(input_dir: pathlib.Path | str | None = None, model_path: pathlib.Path | str | None = None)

See run()

_load_pickle(model_path: pathlib.Path | str)

Helper function to wrap loading a pickle file from a given path for easier testing.

Parameters:

model_path (str or Path) – The file path to the pickle file.

Returns:

The object loaded from the pickle file.

Return type:

object

_transform_batch(batch_tuple: tuple)

Private helper to transform a single batch

Parameters:

batch_tuple (tuple()) – first element is the IDs of the batch as a numpy array second element is the inference results to transform as a numpy array with shape (batch_len, N) where N is the total number of dimensions in the inference result. Caller flattens all inference result axes for us.

Returns:

first element is the ids of the batch as a numpy array second element is the results of running the umap transform on the input as a numpy array.

Return type:

tuple

static _log_memory_usage(message: str = '')

Log the current resident set size (RSS) memory usage of the current process in gigabytes.

Parameters:

message (str, optional) – A descriptive message to include in the log output for context.

Notes

This method is intended for debugging and performance monitoring.

class Infer(config)

Bases: hyrax.verbs.verb_registry.Verb

Inference verb

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'infer'

add_parser_kwargs

description = 'Run inference on a model using a dataset.'

REQUIRED_DATA_GROUPS = ('infer',)

OPTIONAL_DATA_GROUPS = ()

static setup_parser(parser)

We don’t need any parser setup for CLI opts

run_cli(args=None)

CLI stub for Infer verb

run()

Run inference on a model using a dataset

Parameters:

config (dict) – The parsed config file as a nested dict

class Train(config)

Bases: hyrax.verbs.verb_registry.Verb

Train verb

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'train'

add_parser_kwargs

description = 'Train a model using provided data.'

REQUIRED_DATA_GROUPS = ('train',)

OPTIONAL_DATA_GROUPS = ('validate', 'test')

static setup_parser(parser)

We don’t need any parser setup for CLI opts

run_cli(args=None)

CLI stub for Train verb

run()

Run the training process for the configured model and data loader. Returns the trained model.

static _log_params(config, results_dir)

Log the various parameters to mlflow from the config file.

Parameters:

• config (dict) – The main configuration dictionary

• results_dir (str) – The full path to the results sub-directory

class Test(config)

Bases: hyrax.verbs.verb_registry.Verb

Test verb - evaluates a trained model on test data

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'test'

add_parser_kwargs

description = 'Evaluate a trained model on test data.'

REQUIRED_DATA_GROUPS = ('test',)

OPTIONAL_DATA_GROUPS = ()

static setup_parser(parser)

We don’t need any parser setup for CLI opts

run_cli(args=None)

CLI stub for Test verb

run()

Run the test process for the configured model on test data. This evaluates a trained model, saves outputs, and returns metrics.

Note: The configuration dictionary will be updated with the full path to the model weights file that is loaded into the model (config[“test”][“model_weights_file”]).

Returns:

Dataset containing test results that can be used for further analysis

Return type:

InferenceDataset

static _log_params(config, results_dir)

Log the various parameters to mlflow from the config file.

Parameters:

• config (dict) – The main configuration dictionary

• results_dir (str) – The full path to the results sub-directory

class Visualize(config)

Bases: hyrax.verbs.verb_registry.Verb

Verb to create a visualization

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'visualize'

add_parser_kwargs

description = 'Generate a visualization of a latent space created by a UMAP reduction.'

REQUIRED_DATA_GROUPS = ('infer',)

OPTIONAL_DATA_GROUPS = ()

static setup_parser(parser: argparse.ArgumentParser)

CLI not implemented for this verb

run_cli(args: argparse.Namespace | None = None)

CLI not implemented for this verb

run(input_dir: pathlib.Path | str | None = None, *, return_verb: bool = False, make_lupton_rgb_opts: dict | None = None, **kwargs)

Generate an interactive notebook visualization of a latent space that has been umapped down to 2d.

The plot contains two holoviews objects, a scatter plot of the latent space, and a table of objects which can be populated by selecting from the scatter plot.

Parameters:

• input_dir (Optional[Union[Path, str]], optional) – Directory holding the output from the ‘umap’ verb, by default None. When not provided, we use [results][inference_dir] from config. If that’s false; we the most recent umap in the current results directory.

• return_verb (bool, optional) – If True, also return the underlying Visualize instance for post-hoc access to selection state. Defaults to False.

• make_lupton_rgb_opts (dict, optional) – Dictionary of options to pass to astropy’s make_lupton_rgb function for RGB image creation. Default is {“stretch”: 5, “Q”: 8}. Common parameters include stretch (brightness/contrast) and Q (softening parameter for asinh transformation).

• kwargs – Keyword arguments are passed through as options for the plot object as plot_pane.opts(**plot_options). It is not recommended to override the “tools” plot option, because that will break the integration between the plot selection operations and the table.

Returns:

• Holoviews, if return_verb = True (defaul) – A Collection of Haloviews Panes

• tuple of (pane, Visualize), if return_verb = True – Returns a 2-tuple with the pane and the verb instance.

visible_points(x_range: tuple | list, y_range: tuple | list)

Generate a hv.Points object with the points inside the bounding box passed.

This is the event handler for moving or scaling the latent space plot, and is called by Holoviews.

Parameters:

• x_range (tuple or list) – min and max x values

• y_range (tuple or list) – min and max y values

Returns:

Points lying inside the bounding box passed

Return type:

hv.Points

update_points(**kwargs) → None

This is the main UI event handler for selection tools on the plot. If you are a dynamic map in the layout of the visualizer who updates based on plot selection you MUST call this function.

This function accepts the data values from all streams and uses the differences between the current call and prior calls to differentiate between different UI events.

The self.prev_kwargs dictionary is used to store previous calls to this function, and the _called_* helpers perform the differencing for each case.

Calling this function GUARANTEES that self.points, self.points_id, and self.points_idx are up-to-date with the user’s latest selection, regardless of the order that Holoviews evaluates the DynamicMaps in.

_called_lasso(kwargs)

_called_tap(kwargs)

_called_box_select(kwargs)

poly_select_points(geometry) → tuple[numpy.typing.ArrayLike, numpy.typing.ArrayLike, numpy.typing.ArrayLike]

Select points inside a polygon.

Parameters:

geometry (list) – List of x/y points describing the verticies of the polygon

Returns:

First element is an ndarray of x/y points in latent space inside the polygon Second element is an ndarray of corresponding object ids

Return type:

Tuple

box_select_points(x_range: tuple | list, y_range: tuple | list) → tuple[numpy.typing.ArrayLike, numpy.typing.ArrayLike, numpy.typing.ArrayLike]

Return the points and IDs for a box in the latent space

Parameters:

• x_range (tuple or list) – min and max x values

• y_range (tuple or list) – min and max y values

Returns:

First element is an ndarray of x/y points in latent space inside the box Second element is an ndarray of corresponding object ids

Return type:

Tuple

box_select_indexes(x_range: tuple | list, y_range: tuple | list)

Return the indexes inside of a particular box in the latent space

Parameters:

• x_range (tuple or list) – min and max x values

• y_range (tuple or list) – min and max y values

Returns:

Array of data indexes where the latent space representation falls inside the given box.

Return type:

np.ndarray

selected_objects(**kwargs)

Generate the holoview table for a selected set of objects based on input from the Lasso, Tap, and SelectionXY streams.

Returns:

Table with Object ID, x, y locations of the selected objects

Return type:

hv.Table

_table_from_points()

static _bounding_box(points)

_even_aspect_bounding_box()

get_selected_df()

Retrieve a pandas DataFrame containing the currently selected points and their associated metadata.

Returns:

A DataFrame with one row per selected point and columns: [“object_id”, “x”, “y”, *additional_fields].

Return type:

pd.DataFrame

_load_images(**kwargs)

_make_image_pane(total_width: int = 500, *args, **kwargs)

Sample up to 6 of the selected object_ids, load their FITS cutouts from [general][data_dir], and render as small hv.Image thumbnails in a grid.

class VisualizeV2(config)

Bases: hyrax.verbs.verb_registry.Verb

Verb to create a hexbin visualization of a 2D latent space.

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'visualize_v2'

add_parser_kwargs

REQUIRED_DATA_GROUPS = ('visualize',)

OPTIONAL_DATA_GROUPS = ()

static setup_parser(parser: argparse.ArgumentParser)

CLI not implemented for this verb

run_cli(args: argparse.Namespace | None = None)

CLI not implemented for this verb

run(**kwargs)

Generate an interactive hexbin visualization of a latent space projected to 2D.

Uses HoloViews HexTiles with datashader for adaptive hexbin aggregation, box/lasso selection, a metadata table, and tabbed detail plots.

Parameters:

kwargs – Additional keyword arguments passed as HexTiles opts overrides.

Returns:

This verb instance. Use it to call restart_ui() or get_selected_df() after the UI has been displayed.

Return type:

VisualizeV2

restart_ui(**kwargs)

Rebuild and re-display the Panel UI without reloading data.

Call this after a Jupyter websocket disconnect instead of re-running the cell. The expensive data-loading step is skipped — only the widgets are rebuilt.

Parameters:

kwargs – Additional keyword arguments passed as HexTiles opts overrides.

Returns:

This verb instance. Use it to call restart_ui() or get_selected_df() after the UI has been displayed.

Return type:

VisualizeV2

_load_data()

Load dataset and build the points DataFrame.

Guards with a _data_loaded sentinel so the expensive steps only run once per verb instance. Safe to call multiple times.

_build_ui(**kwargs)

Build and display the Panel UI using data already loaded by _load_data().

get_selected_df()

Return the current selection as a DataFrame.

class Lookup(config)

Bases: hyrax.verbs.verb_registry.Verb

Look up an inference result using the ID of a data member

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'lookup'

add_parser_kwargs

description = 'Look up an inference result using the ID of a data member.'

static setup_parser(parser: argparse.ArgumentParser)

Set up our arguments by configuring a subparser

Parameters:

parser (ArgumentParser) – The sub-parser to configure

run_cli(args: argparse.Namespace | None = None)

Entrypoint to Lookup from the CLI.

Parameters:

args (Optional[Namespace], optional) – The parsed command line arguments

run(id: str, results_dir: pathlib.Path | str | None = None) → numpy.ndarray | None

Lookup the latent-space representation of a particular ID

Requires the relevant dataset to be configured, and for inference to have been run.

Parameters:

• id (str) – The ID of the input data to look up the inference result

• results_dir (str, Optional) – The directory containing the inference results.

Returns:

The output tensor of the model for the given input.

Return type:

Optional[np.ndarray]

class SaveToDatabase(config)

Bases: hyrax.verbs.verb_registry.Verb

Verb to insert inference results into a vector database index for fast similarity search.

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'save_to_database'

add_parser_kwargs

description = 'Insert inference results into vector database.'

static setup_parser(parser: argparse.ArgumentParser)

Stub of parser setup

run_cli(args: argparse.Namespace | None = None)

Stub CLI implementation

run(input_dir: pathlib.Path | str | None = None, output_dir: pathlib.Path | str | None = None)

Insert inference results into vector database.

Parameters:

• input_dir (str or Path, Optional) – The directory containing the inference results.

• output_dir (str or Path, Optional) – The directory where the vector database is stored. If None, a new directory will be created. If specified, it can point to either an empty directory or a directory containing an existing vector database. If the latter, the database will be updated with the new vectors.

class Model(config)

Bases: hyrax.verbs.verb_registry.Verb

Resolves the model class that is defined in the config file. This will return a reference to the model class.

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'model'

add_parser_kwargs

description = 'Return a reference to the model class (not a new instance).'

static setup_parser(parser)

Not implemented

run_cli()

Not implemented

run()

Fetch and return the model _class_. Does not create an instance of the model class.

class ToOnnx(config)

Bases: hyrax.verbs.verb_registry.Verb

Export the model to ONNX format

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'to_onnx'

add_parser_kwargs

description = 'Export model to ONNX format.'

static setup_parser(parser)

Setup parser for ONNX export verb

run_cli(args=None)

Run the ONNX export verb from the CLI

run(input_model_directory: str = None)

Export the model to ONNX format and save it to the specified path.

class Engine(config)

Bases: hyrax.verbs.verb_registry.Verb

This verb drives inference with an ONNX model in production.

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'engine'

add_parser_kwargs

description = 'Run inference with an ONNX model.'

static setup_parser(parser)

Setup parser for engine verb

run_cli(args=None)

CLI stub for Engine verb

run(model_directory: str = None)

Run inference with an ONNX model.

This method performs the following steps: - Read in the user config - Prepare all the datasets requested - Implement a simple strategy for reading in batches of data samples - Process the samples with any custom collate functions as well as a default collate function - Pass the collated batch to the appropriate to_tensor function - Send that output to the ONNX-ified model - Persist the results of inference

Parameters:

model_directory (str, optional) – Directory containing the ONNX model. If not provided, uses the config file or finds the most recent ONNX export directory.

create_ort_inputs(prepared_batch)

Create the inputs array for the ONNX model using the expected inputs from the loaded ONNX model and the type and shape of the prepared batch.

run_onnx_batch(ort_inputs)

Run the batch using our onnx runtime session

Only split out because this is when data is mutated and we need to be able to trace it.

_setup_trace(prepare_inputs_fn)

class Prepare(config)

Bases: hyrax.verbs.verb_registry.Verb

Prepare Verb, Prepares a dataset and returns it

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'prepare'

add_parser_kwargs

static setup_parser(parser)

We don’t need any parser setup for CLI opts

run_cli(args=None)

CLI stub for Prepare verb

run()

Prepare the dataset for a given model and data loader using the verb’s configuration.

Uses self.config to construct and return the prepared dataset.

class Verb(config)

Bases: abc.ABC

Base class for all hyrax verbs

__init__()

Overall initialization for all verbs that saves the config

add_parser_kwargs: dict[str, str]

REQUIRED_DATA_GROUPS: tuple[str, Ellipsis] = ()

OPTIONAL_DATA_GROUPS: tuple[str, Ellipsis] = ()

cli_name = 'VERB'

description = ''

config

classmethod information()

Returns a string describing this verb. Includes the following: - Name of the verb - Required Data Groups - Optional Data Groups - One line description of what this verb does

If a data group is empty then it will be printed as an empty tuple.

Returns:

<name>: Data Groups: Req. (<req1>, <req2>, …), Opt. (<opt1>, <opt2>, …). <Description>

Return type:

str

validate_data_request() → None

Validate the data_request configuration for this verb’s known groups.

Reads data_request from the verb’s config and checks:

1. All groups listed in REQUIRED_DATA_GROUPS are present.

2. Cross-group split_fraction constraints (sum ≤ 1.0, consistency) hold for the active groups only — groups outside REQUIRED_DATA_GROUPS + OPTIONAL_DATA_GROUPS are ignored so that unrelated groups in a shared config do not cause false failures.

Verbs that define neither REQUIRED_DATA_GROUPS nor OPTIONAL_DATA_GROUPS skip validation entirely.

Raises:

RuntimeError – If a required group is absent, or if cross-group split_fraction constraints are violated for the active groups.

all_class_verbs() → list[str]

Returns all verbs that are currently registered with a class-based implementation

all_verbs() → list[str]

Returns all verbs that are currently registered

fetch_verb_class(cli_name: str) → type[Verb] | None

Gives the class object for the named verb

Parameters:

cli_name (str) – The name of the verb on the command line interface

Returns:

The verb class or None if no such verb class exists.

Return type:

Optional[type[Verb]]

is_verb_class(cli_name: str) → bool

Returns true if the verb has a class based implementation

Parameters:

cli_name (str) – The name of the verb on the command line interface

Returns:

True if the verb has a class-based implementation

Return type:

bool