hyrax.verbs

Submodules

• hyrax.verbs.database_connection
• hyrax.verbs.engine
• hyrax.verbs.infer
• hyrax.verbs.lookup
• hyrax.verbs.model
• hyrax.verbs.save_to_database
• hyrax.verbs.search
• hyrax.verbs.to_onnx
• hyrax.verbs.train
• hyrax.verbs.umap
• hyrax.verbs.verb_registry
• hyrax.verbs.visualize

Classes

DatabaseConnection	Verb to insert inference results into a vector database index for fast
Umap	Umap latent space points into 2d
Infer	Inference verb
Train	Train verb
Visualize	Verb to create a visualization
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.
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 insert inference results into a vector database index for fast similarity search.

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'database_connection'

add_parser_kwargs

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

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)

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.

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)

See run()

_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

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

static load_model_weights(config, model)

Loads the model weights from a file. Raises RuntimeError if this is not possible due to config, missing or malformed file

Parameters:

• config (dict) – Full runtime configuration

• model (nn.Module) – The model class to load weights into

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

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 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

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 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

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

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

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

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

static setup_parser(parser)

Setup parser for engine verb

run_cli(args=None)

CLI stub for Engine verb

run(model_directory: str = None)

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

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]

config

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