hyrax.verbs =========== .. py:module:: hyrax.verbs Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/hyrax/verbs/database_connection/index /autoapi/hyrax/verbs/engine/index /autoapi/hyrax/verbs/infer/index /autoapi/hyrax/verbs/lookup/index /autoapi/hyrax/verbs/model/index /autoapi/hyrax/verbs/prepare/index /autoapi/hyrax/verbs/save_to_database/index /autoapi/hyrax/verbs/test/index /autoapi/hyrax/verbs/to_onnx/index /autoapi/hyrax/verbs/train/index /autoapi/hyrax/verbs/umap/index /autoapi/hyrax/verbs/verb_registry/index /autoapi/hyrax/verbs/visualize/index Classes ------- .. autoapisummary:: hyrax.verbs.DatabaseConnection hyrax.verbs.Umap hyrax.verbs.Infer hyrax.verbs.Train hyrax.verbs.Test hyrax.verbs.Visualize hyrax.verbs.Lookup hyrax.verbs.SaveToDatabase hyrax.verbs.Model hyrax.verbs.ToOnnx hyrax.verbs.Engine hyrax.verbs.Prepare hyrax.verbs.Verb Functions --------- .. autoapisummary:: hyrax.verbs.all_class_verbs hyrax.verbs.all_verbs hyrax.verbs.fetch_verb_class hyrax.verbs.is_verb_class Package Contents ---------------- .. py:class:: DatabaseConnection(config) Bases: :py:obj:`hyrax.verbs.verb_registry.Verb` Verb to create a connection to a vector database with inference results. .. py:method:: __init__ Overall initialization for all verbs that saves the config .. py:attribute:: cli_name :value: 'database_connection' .. py:attribute:: add_parser_kwargs .. py:attribute:: description :value: 'Create a connection to the vector database for interactive queries.' .. py:method:: setup_parser(parser: argparse.ArgumentParser) :staticmethod: Stub of parser setup .. py:method:: run_cli(args: argparse.Namespace | None = None) Stub CLI implementation .. py:method:: run(database_dir: Union[pathlib.Path, str] | None = None) Create a connection to the vector database for interactive queries. :param database_dir: 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. :type database_dir: str or Path, Optional .. py:method:: _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". :param database_dir: The directory containing the vector database and the config file that be used as reference. :type database_dir: Path :returns: The config value for ["vector_db"]["name"] in the reference config. :rtype: str .. py:class:: Umap(config) Bases: :py:obj:`hyrax.verbs.verb_registry.Verb` Umap latent space points into 2d .. py:method:: __init__ Overall initialization for all verbs that saves the config .. py:attribute:: cli_name :value: 'umap' .. py:attribute:: add_parser_kwargs .. py:attribute:: description :value: 'Transforms the entire dataset into a lower-dimensional space by fitting a UMAP model.' .. py:method:: setup_parser(parser: argparse.ArgumentParser) :staticmethod: Stub of parser setup .. py:method:: run_cli(args: argparse.Namespace | None = None) Stub CLI implementation .. py:method:: run(input_dir: Union[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. :param input_dir: The directory containing the inference results. :type input_dir: str or Path, Optional :returns: The method does not return anything but saves the UMAP representations to disk. :rtype: None .. py:method:: _run(input_dir: Union[pathlib.Path, str] | None = None) See run() .. py:method:: _transform_batch(batch_tuple: tuple) Private helper to transform a single batch :param batch_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. :type batch_tuple: tuple() :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. :rtype: tuple .. py:method:: _log_memory_usage(message: str = '') :staticmethod: Log the current resident set size (RSS) memory usage of the current process in gigabytes. :param message: A descriptive message to include in the log output for context. :type message: str, optional .. rubric:: Notes This method is intended for debugging and performance monitoring. .. py:class:: Infer(config) Bases: :py:obj:`hyrax.verbs.verb_registry.Verb` Inference verb .. py:method:: __init__ Overall initialization for all verbs that saves the config .. py:attribute:: cli_name :value: 'infer' .. py:attribute:: add_parser_kwargs .. py:attribute:: description :value: 'Run inference on a model using a dataset.' .. py:attribute:: REQUIRED_DATA_GROUPS :value: ('infer',) .. py:attribute:: OPTIONAL_DATA_GROUPS :value: () .. py:method:: setup_parser(parser) :staticmethod: We don't need any parser setup for CLI opts .. py:method:: run_cli(args=None) CLI stub for Infer verb .. py:method:: run() Run inference on a model using a dataset :param config: The parsed config file as a nested dict :type config: dict .. py:class:: Train(config) Bases: :py:obj:`hyrax.verbs.verb_registry.Verb` Train verb .. py:method:: __init__ Overall initialization for all verbs that saves the config .. py:attribute:: cli_name :value: 'train' .. py:attribute:: add_parser_kwargs .. py:attribute:: description :value: 'Train a model using provided data.' .. py:attribute:: REQUIRED_DATA_GROUPS :value: ('train',) .. py:attribute:: OPTIONAL_DATA_GROUPS :value: ('validate', 'test') .. py:method:: setup_parser(parser) :staticmethod: We don't need any parser setup for CLI opts .. py:method:: run_cli(args=None) CLI stub for Train verb .. py:method:: run() Run the training process for the configured model and data loader. Returns the trained model. .. py:method:: _log_params(config, results_dir) :staticmethod: Log the various parameters to mlflow from the config file. :param config: The main configuration dictionary :type config: dict :param results_dir: The full path to the results sub-directory :type results_dir: str .. py:class:: Test(config) Bases: :py:obj:`hyrax.verbs.verb_registry.Verb` Test verb - evaluates a trained model on test data .. py:method:: __init__ Overall initialization for all verbs that saves the config .. py:attribute:: cli_name :value: 'test' .. py:attribute:: add_parser_kwargs .. py:attribute:: description :value: 'Evaluate a trained model on test data.' .. py:attribute:: REQUIRED_DATA_GROUPS :value: ('test',) .. py:attribute:: OPTIONAL_DATA_GROUPS :value: () .. py:method:: setup_parser(parser) :staticmethod: We don't need any parser setup for CLI opts .. py:method:: run_cli(args=None) CLI stub for Test verb .. py:method:: 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 :rtype: InferenceDataset .. py:method:: _log_params(config, results_dir) :staticmethod: Log the various parameters to mlflow from the config file. :param config: The main configuration dictionary :type config: dict :param results_dir: The full path to the results sub-directory :type results_dir: str .. py:class:: Visualize(config) Bases: :py:obj:`hyrax.verbs.verb_registry.Verb` Verb to create a visualization .. py:method:: __init__ Overall initialization for all verbs that saves the config .. py:attribute:: cli_name :value: 'visualize' .. py:attribute:: add_parser_kwargs .. py:attribute:: description :value: 'Generate a visualization of a latent space created by a UMAP reduction.' .. py:attribute:: REQUIRED_DATA_GROUPS :value: ('infer',) .. py:attribute:: OPTIONAL_DATA_GROUPS :value: () .. py:method:: setup_parser(parser: argparse.ArgumentParser) :staticmethod: CLI not implemented for this verb .. py:method:: run_cli(args: argparse.Namespace | None = None) CLI not implemented for this verb .. py:method:: run(input_dir: Union[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. :param input_dir: 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. :type input_dir: Optional[Union[Path, str]], optional :param return_verb: If True, also return the underlying Visualize instance for post-hoc access to selection state. Defaults to False. :type return_verb: bool, optional :param make_lupton_rgb_opts: 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). :type make_lupton_rgb_opts: dict, optional :param 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. .. py:method:: visible_points(x_range: Union[tuple, list], y_range: Union[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. :param x_range: min and max x values :type x_range: tuple or list :param y_range: min and max y values :type y_range: tuple or list :returns: Points lying inside the bounding box passed :rtype: hv.Points .. py:method:: 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. .. py:method:: _called_lasso(kwargs) .. py:method:: _called_tap(kwargs) .. py:method:: _called_box_select(kwargs) .. py:method:: poly_select_points(geometry) -> tuple[numpy.typing.ArrayLike, numpy.typing.ArrayLike, numpy.typing.ArrayLike] Select points inside a polygon. :param geometry: List of x/y points describing the verticies of the polygon :type geometry: list :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 :rtype: Tuple .. py:method:: box_select_points(x_range: Union[tuple, list], y_range: Union[tuple, list]) -> tuple[numpy.typing.ArrayLike, numpy.typing.ArrayLike, numpy.typing.ArrayLike] Return the points and IDs for a box in the latent space :param x_range: min and max x values :type x_range: tuple or list :param y_range: min and max y values :type y_range: tuple or list :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 :rtype: Tuple .. py:method:: box_select_indexes(x_range: Union[tuple, list], y_range: Union[tuple, list]) Return the indexes inside of a particular box in the latent space :param x_range: min and max x values :type x_range: tuple or list :param y_range: min and max y values :type y_range: tuple or list :returns: Array of data indexes where the latent space representation falls inside the given box. :rtype: np.ndarray .. py:method:: 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 :rtype: hv.Table .. py:method:: _table_from_points() .. py:method:: _bounding_box(points) :staticmethod: .. py:method:: _even_aspect_bounding_box() .. py:method:: 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]. :rtype: pd.DataFrame .. py:method:: _load_images(**kwargs) .. py:method:: _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. .. py:class:: Lookup(config) Bases: :py:obj:`hyrax.verbs.verb_registry.Verb` Look up an inference result using the ID of a data member .. py:method:: __init__ Overall initialization for all verbs that saves the config .. py:attribute:: cli_name :value: 'lookup' .. py:attribute:: add_parser_kwargs .. py:attribute:: description :value: 'Look up an inference result using the ID of a data member.' .. py:method:: setup_parser(parser: argparse.ArgumentParser) :staticmethod: Set up our arguments by configuring a subparser :param parser: The sub-parser to configure :type parser: ArgumentParser .. py:method:: run_cli(args: argparse.Namespace | None = None) Entrypoint to Lookup from the CLI. :param args: The parsed command line arguments :type args: Optional[Namespace], optional .. py:method:: run(id: str, results_dir: Union[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. :param id: The ID of the input data to look up the inference result :type id: str :param results_dir: The directory containing the inference results. :type results_dir: str, Optional :returns: The output tensor of the model for the given input. :rtype: Optional[np.ndarray] .. py:class:: SaveToDatabase(config) Bases: :py:obj:`hyrax.verbs.verb_registry.Verb` Verb to insert inference results into a vector database index for fast similarity search. .. py:method:: __init__ Overall initialization for all verbs that saves the config .. py:attribute:: cli_name :value: 'save_to_database' .. py:attribute:: add_parser_kwargs .. py:attribute:: description :value: 'Insert inference results into vector database.' .. py:method:: setup_parser(parser: argparse.ArgumentParser) :staticmethod: Stub of parser setup .. py:method:: run_cli(args: argparse.Namespace | None = None) Stub CLI implementation .. py:method:: run(input_dir: Union[pathlib.Path, str] | None = None, output_dir: Union[pathlib.Path, str] | None = None) Insert inference results into vector database. :param input_dir: The directory containing the inference results. :type input_dir: str or Path, Optional :param output_dir: 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. :type output_dir: str or Path, Optional .. py:class:: Model(config) Bases: :py:obj:`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. .. py:method:: __init__ Overall initialization for all verbs that saves the config .. py:attribute:: cli_name :value: 'model' .. py:attribute:: add_parser_kwargs .. py:attribute:: description :value: 'Return a reference to the model class (not a new instance).' .. py:method:: setup_parser(parser) :staticmethod: Not implemented .. py:method:: run_cli() Not implemented .. py:method:: run() Fetch and return the model _class_. Does not create an instance of the model class. .. py:class:: ToOnnx(config) Bases: :py:obj:`hyrax.verbs.verb_registry.Verb` Export the model to ONNX format .. py:method:: __init__ Overall initialization for all verbs that saves the config .. py:attribute:: cli_name :value: 'to_onnx' .. py:attribute:: add_parser_kwargs .. py:attribute:: description :value: 'Export model to ONNX format.' .. py:method:: setup_parser(parser) :staticmethod: Setup parser for ONNX export verb .. py:method:: run_cli(args=None) Run the ONNX export verb from the CLI .. py:method:: run(input_model_directory: str = None) Export the model to ONNX format and save it to the specified path. .. py:class:: Engine(config) Bases: :py:obj:`hyrax.verbs.verb_registry.Verb` This verb drives inference with an ONNX model in production. .. py:method:: __init__ Overall initialization for all verbs that saves the config .. py:attribute:: cli_name :value: 'engine' .. py:attribute:: add_parser_kwargs .. py:attribute:: description :value: 'Run inference with an ONNX model.' .. py:method:: setup_parser(parser) :staticmethod: Setup parser for engine verb .. py:method:: run_cli(args=None) CLI stub for Engine verb .. py:method:: 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 :param model_directory: Directory containing the ONNX model. If not provided, uses the config file or finds the most recent ONNX export directory. :type model_directory: str, optional .. py:method:: 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. .. py:method:: 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. .. py:method:: _setup_trace(prepare_inputs_fn) .. py:class:: Prepare(config) Bases: :py:obj:`hyrax.verbs.verb_registry.Verb` Prepare Verb, Prepares a dataset and returns it .. py:method:: __init__ Overall initialization for all verbs that saves the config .. py:attribute:: cli_name :value: 'prepare' .. py:attribute:: add_parser_kwargs .. py:method:: setup_parser(parser) :staticmethod: We don't need any parser setup for CLI opts .. py:method:: run_cli(args=None) CLI stub for Prepare verb .. py:method:: 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. .. py:class:: Verb(config) Bases: :py:obj:`abc.ABC` Base class for all hyrax verbs .. py:method:: __init__ Overall initialization for all verbs that saves the config .. py:attribute:: add_parser_kwargs :type: dict[str, str] .. py:attribute:: REQUIRED_DATA_GROUPS :type: tuple[str, Ellipsis] :value: () .. py:attribute:: OPTIONAL_DATA_GROUPS :type: tuple[str, Ellipsis] :value: () .. py:attribute:: cli_name :value: 'VERB' .. py:attribute:: description :value: '' .. py:attribute:: config .. py:method:: information() :classmethod: 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: : Data Groups: Req. (, , ...), Opt. (, , ...). :rtype: str .. py:method:: validate_data_request() -> None Validate the data_request configuration for this verb's known groups. Reads ``data_request`` (or the deprecated ``model_inputs``) 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. .. py:function:: all_class_verbs() -> list[str] Returns all verbs that are currently registered with a class-based implementation .. py:function:: all_verbs() -> list[str] Returns all verbs that are currently registered .. py:function:: fetch_verb_class(cli_name: str) -> type[Verb] | None Gives the class object for the named verb :param cli_name: The name of the verb on the command line interface :type cli_name: str :returns: The verb class or None if no such verb class exists. :rtype: Optional[type[Verb]] .. py:function:: is_verb_class(cli_name: str) -> bool Returns true if the verb has a class based implementation :param cli_name: The name of the verb on the command line interface :type cli_name: str :returns: True if the verb has a class-based implementation :rtype: bool