Utils

float_time_to_index(time, num_time_steps)

Convert a float time value to a time index.

Parameters:

Name	Type	Description	Default
time	Tensor	A tensor of float time values in the range [0, 1].	required
num_time_steps	int	The number of discrete time steps.	required

Returns:

Type	Description
Tensor	torch.Tensor: A tensor of time indices corresponding to the input float time values.

Source code in bionemo/moco/distributions/time/utils.py

def float_time_to_index(time: torch.Tensor, num_time_steps: int) -> torch.Tensor:
    """Convert a float time value to a time index.

    Args:
        time (torch.Tensor): A tensor of float time values in the range [0, 1].
        num_time_steps (int): The number of discrete time steps.

    Returns:
        torch.Tensor: A tensor of time indices corresponding to the input float time values.
    """
    # Ensure time values are in the range [0, 1]
    time = torch.clamp(time, 0.0, 1.0)

    # Scale to the index range and round
    indices = torch.round(time * (num_time_steps - 1)).to(torch.int64)

    return indices