Update docstrings and add missing types

pyrad/cameras/cameras.py

@@ -187,22 +187,23 @@ class PinholeCamera(Camera):
         return torch.tensor([self.cx, self.cy, self.fx, self.fy])
 
     @classmethod
-    def fx_index(cls):
-        """TODO(ethan): redo this in a better way.
+    def fx_index(cls) -> int:
+        """Returns the index of the fx parameter in the intrinsics vector
+        TODO(ethan): redo this in a better way.
         Ideally we can dynamically grab the focal length parameters depending on
         Simple vs. not Simple Pinhole Model.
 
         Returns:
-            _type_: _description_
+            x focal length index
         """
         return 2
 
     @classmethod
-    def fy_index(cls):
-        """_summary_
+    def fy_index(cls) -> int:
+        """Returns the index of the fy parameter in the intrinsics vector
 
         Returns:
-            _type_: _description_
+            y focal length index
         """
         return 3
 

pyrad/engine/trainer.py

@@ -67,7 +67,11 @@ class Trainer:
         self.visualizer_state = viewer_utils.VisualizerState(config.viewer)
 
     def setup(self, test_mode=False):
-        """Setup the Trainer by calling other setup functions."""
+        """Setup the Trainer by calling other setup functions.
+
+        Args:
+            test_mode (bool, optional): Whether to setup for testing. Defaults to False.
+        """
         self.dataset_inputs_train, self.dataloader_train = setup_dataset_train(self.config.data, device=self.device)
         _, self.dataloader_eval = setup_dataset_eval(self.config.data, test_mode=test_mode, device=self.device)
         self.graph = setup_graph(self.config.graph, self.dataset_inputs_train, device=self.device)
@@ -134,7 +138,7 @@ class Trainer:
             writer.write_out_storage()
 
     def _load_checkpoint(self) -> None:
-        """helper function to load graph and optimizer from prespecified checkpoint"""
+        """Helper function to load graph and optimizer from prespecified checkpoint"""
         load_config = self.config.trainer.resume_train
         load_path = os.path.join(load_config.load_dir, f"step-{load_config.load_step:09d}.ckpt")
         assert os.path.exists(load_path), f"Checkpoint {load_path} does not exist"
@@ -150,8 +154,8 @@ class Trainer:
         """Save the model and optimizers
 
         Args:
-            output_dir (str): directory to save the checkpoint
-            step (int): number of steps in training for given checkpoint
+            output_dir: directory to save the checkpoint
+            step: number of steps in training for given checkpoint
         """
         if not os.path.exists(output_dir):
             os.makedirs(output_dir)
@@ -170,9 +174,9 @@ class Trainer:
         """Run one iteration with a batch of inputs.
 
         Args:
-            ray_indices (TensorType["num_rays", 3]): Contains camera, row, and col indicies for target rays.
-            batch (dict): Batch of training data.
-            step (int): Current training step.
+            ray_indices: Contains camera, row, and col indicies for target rays.
+            batch: Batch of training data.
+            step: Current training step.
 
         Returns:
             Dict[str, float]: Dictionary of model losses.
@@ -193,9 +197,9 @@ class Trainer:
         """Test a specific image.
 
         Args:
-            camera_ray_bundle (RayBundle): Bundle of test rays.
-            batch (dict): Batch of data.
-            step (int): Current training step.
+            camera_ray_bundle: Bundle of test rays.
+            batch: Batch of data.
+            step: Current training step.
 
         Returns:
             float: PSNR
@@ -211,8 +215,8 @@ class Trainer:
         """Run evaluation with a given dataloader.
 
         Args:
-            dataloader (EvalDataLoader): Evaluation dataloader.
-            step (int): Current training iteration.
+            dataloader: Evaluation dataloader.
+            step: Current training iteration.
         """
         for camera_ray_bundle, batch in dataloader:
             self.test_image(camera_ray_bundle, batch, step=step)

pyrad/optimizers/optimizers.py

@@ -31,11 +31,11 @@ def setup_optimizers(config: DictConfig, param_groups: Dict[str, List[Parameter]
     """Helper to set up the optimizers
 
     Args:
-        config (DictConfig): _description_
-        param_groups (Dict[str, List[Parameter]]): _description_
+        config: The optimizer configuration object.
+        param_groups: A dictionary of parameter groups to optimize.
 
     Returns:
-        Optimizers: _description_
+        The optimizers object.
     """
     return Optimizers(config, param_groups)
 
@@ -45,6 +45,14 @@ class ExponentialDecaySchedule(LambdaLR):
     See https://github.com/google-research/google-research/blob/
     fd2cea8cdd86b3ed2c640cbe5561707639e682f3/jaxnerf/nerf/utils.py#L360
     for details.
+
+    Args:
+        optimizer: The optimizer to update.
+        lr_init: The initial learning rate.
+        lr_final: The final learning rate.
+        max_steps: The maximum number of steps.
+        lr_delay_steps: The number of steps to delay the learning rate. Default is 0.
+        lr_delay_mult: The multiplier for the learning rate after the delay. Default is 1.0.
     """
 
     def __init__(self, optimizer, lr_init, lr_final, max_steps, lr_delay_steps=0, lr_delay_mult=1.0) -> None:
@@ -66,15 +74,14 @@ class ExponentialDecaySchedule(LambdaLR):
 
 
 class Optimizers:
-    """_summary_"""
+    """A set of optimizers.
+
+    Args:
+        config: The optimizer configuration object.
+        param_groups: A dictionary of parameter groups to optimize.
+    """
 
     def __init__(self, config: DictConfig, param_groups: Dict[str, List[Parameter]]):
-        """_summary_
-
-        Args:
-            config (DictConfig): _description_
-            param_dict (Dict[str, List[Parameter]]): _description_
-        """
         self.config = config
         self.optimizers = {}
         self.schedulers = {}
@@ -87,35 +94,39 @@ class Optimizers:
                 )
 
     def optimizer_step(self, param_group_name: str) -> None:
-        """fetch and step corresponding optimizer
+        """Fetch and step corresponding optimizer.
 
         Args:
-            param_group_name (str): name of optimizer to step forward
+            param_group_name: name of optimizer to step forward
         """
         self.optimizers[param_group_name].step()
 
     def scheduler_step(self, param_group_name: str) -> None:
-        """fetch and step corresponding scheduler
+        """Fetch and step corresponding scheduler.
 
         Args:
-            param_group_name (str): name of scheduler to step forward
+            param_group_name: name of scheduler to step forward
         """
         if self.config.param_group_name.scheduler:
             self.schedulers[param_group_name].step()
 
-    def zero_grad_all(self):
-        """zero the gradients for all optimizer parameters"""
+    def zero_grad_all(self) -> None:
+        """Zero the gradients for all optimizer parameters."""
         for _, optimizer in self.optimizers.items():
             optimizer.zero_grad()
 
-    def optimizer_step_all(self):
+    def optimizer_step_all(self) -> None:
         """Run step for all optimizers."""
         for _, optimizer in self.optimizers.items():
             # note that they key is the parameter name
             optimizer.step()
 
-    def scheduler_step_all(self, step):
-        """Run step for all schedulers."""
+    def scheduler_step_all(self, step: int) -> None:
+        """Run step for all schedulers.
+
+        Args:
+            step: the current step
+        """
         for param_group_name, scheduler in self.schedulers.items():
             scheduler.step()
             # TODO(ethan): clean this up. why is there indexing into a list?
@@ -123,6 +134,10 @@ class Optimizers:
             writer.put_scalar(name=f"learning_rate/{param_group_name}", scalar=lr, step=step)
 
     def load_optimizers(self, loaded_state: Dict[str, Any]) -> None:
-        """Helper to load the optimizer state from previous checkpoint"""
+        """Helper to load the optimizer state from previous checkpoint
+
+        Args:
+            loaded_state: the state from the previous checkpoint
+        """
         for k, v in loaded_state["optimizers"].items():
             self.optimizers[k].load_state_dict(v)

pyrad/renderers/renderers.py

@@ -41,7 +41,7 @@ class RGBRenderer(nn.Module):
     """Standard volumetic rendering.
 
     Args:
-        background_color (TensorType[3], optional): Background color as RGB. Defaults to random.
+        background_color: Background color as RGB. Defaults to random.
     """
 
     def __init__(self, background_color: Optional[TensorType[3]] = None) -> None:
@@ -58,12 +58,12 @@ class RGBRenderer(nn.Module):
         """Composite samples along ray and render color image
 
         Args:
-            rgb (TensorType[..., "num_samples", 3]): RGB for each sample
-            weights (TensorType[..., "num_samples", 1]): Weights for each sample
-            background_color (TensorType[3], optional): Background color as RGB. Defaults to random.
+            rgb: RGB for each sample
+            weights: Weights for each sample
+            background_color: Background color as RGB. Defaults to random.
 
         Returns:
-            TensorType[..., 3]: Outputs rgb values.
+            Outputs rgb values.
         """
         rgb = torch.sum(weights * rgb, dim=-2)
 
@@ -82,11 +82,11 @@ class RGBRenderer(nn.Module):
         """Composite samples along ray and render color image
 
         Args:
-            rgb (TensorType[..., "num_samples", 3]): RGB for each sample
-            weights (TensorType[..., "num_samples", 1]): Weights for each sample
+            rgb: RGB for each sample
+            weights: Weights for each sample
 
         Returns:
-            TensorType[..., 3]: Outputs of rgb values.
+            Outputs of rgb values.
         """
 
         rgb = self.combine_rgb(rgb, weights, background_color=self.background_color)
@@ -99,8 +99,8 @@ class SHRenderer(nn.Module):
     """Render RGB value from spherical harmonics.
 
     Args:
-        background_color (TensorType[3], optional): Background color as RGB. Defaults to random.
-        activation (Optional[nn.Module], optional): Output activation. Defaults to Sigmoid().
+        background_color: Background color as RGB. Defaults to random.
+        activation: Output activation. Defaults to Sigmoid().
     """
 
     def __init__(
@@ -119,12 +119,12 @@ class SHRenderer(nn.Module):
         """Composite samples along ray and render color image
 
         Args:
-            sh (TensorType[..., "num_samples", "coeffs"]): Spherical hamonics coefficients for each sample
-            directions: (TensorType[..., "num_samples", 3]): Sample direction
-            weights (TensorType[..., "num_samples", 1]): Weights for each sample
+            sh: Spherical hamonics coefficients for each sample
+            directions: Sample direction
+            weights: Weights for each sample
 
         Returns:
-            TensorType[..., 3]: Outputs of rgb values.
+            Outputs of rgb values.
         """
 
         sh = sh.view(*sh.shape[:-1], 3, sh.shape[-1] // 3)
@@ -154,10 +154,10 @@ class AccumulationRenderer(nn.Module):
         """Composite samples along ray and calculate accumulation.
 
         Args:
-            weights (TensorType[..., "num_samples"]): Weights for each sample
+            weights: Weights for each sample
 
         Returns:
-            TensorType: Outputs of accumulated values.
+            Outputs of accumulated values.
         """
 
         accumulation = torch.sum(weights, dim=-2)
@@ -181,11 +181,11 @@ class DepthRenderer(nn.Module):
         """Composite samples along ray and calculate disparities.
 
         Args:
-            weights (TensorType[..., "num_samples", 1]): Weights for each sample.
-            ray_samples (RaySamples): Set of ray samples.
+            weights: Weights for each sample.
+            ray_samples: Set of ray samples.
 
         Returns:
-            TensorType[..., 1]: Outputs of depth values.
+            Outputs of depth values.
         """
 
         if self.method == "expected":
@@ -207,14 +207,14 @@ class UncertaintyRenderer(nn.Module):
     def forward(
         cls, betas: TensorType[..., "num_samples", 1], weights: TensorType[..., "num_samples", 1]
     ) -> TensorType:
-        """_summary_
+        """Calculate uncertainty along the ray.
 
         Args:
-            betas (TensorType[..., num_samples, 1]): _description_
-            weights (TensorType[..., num_samples, 1]): _description_
+            betas: Uncertainty betas for each sample.
+            weights: Weights of each sample.
 
         Returns:
-            TensorType: _description_
+            Rendering of uncertainty.
         """
         uncertainty = torch.sum(weights * betas, dim=-2)
         return uncertainty

pyrad/utils/io.py

@@ -19,10 +19,15 @@ import json
 import logging
 import os
 import pickle
+from typing import Any, Tuple, Optional
 
 
 def load_from_json(filename: str):
-    """Load a dictionary from a JSON filename."""
+    """Load a dictionary from a JSON filename.
+
+    Args:
+        filename: The filename to load from.
+    """
     assert filename.endswith(".json")
     with open(filename, "r", encoding="UTF-8") as file:
         return json.load(file)
@@ -32,8 +37,8 @@ def write_to_json(filename: str, content: dict):
     """Write data to a JSON file.
 
     Args:
-        filename (str): The filename to write to.
-        content (dict): The dictionary data to write.
+        filename: The filename to write to.
+        content: The dictionary data to write.
     """
     assert filename.endswith(".json")
     with open(filename, "w", encoding="UTF-8") as file:
@@ -41,30 +46,39 @@ def write_to_json(filename: str, content: dict):
 
 
 def load_from_pkl(filename: str):
-    """Load from a pickle file."""
+    """Load from a pickle file.
+
+    Args:
+        filename: The filename to load from.
+    """
     assert filename.endswith(".pkl")
     with open(filename, "rb") as f:
         return pickle.load(f)
 
 
-def write_to_pkl(filename: str, content):
-    "Write to a pickle file."
+def write_to_pkl(filename: str, content: Any):
+    """Write to a pickle file.
+
+    Args:
+        filename (str): The filename to write to.
+        content (Any): The data to write.
+    """
     assert filename.endswith(".pkl")
     with open(filename, "wb") as f:
         pickle.dump(content, f)
 
 
-def get_git_root(path, dirs=(".git",), default=None):
-    """_summary_
-    "https://stackoverflow.com/questions/22081209/find-the-root-of-the-git-repository-where-the-file-lives
+def get_git_root(path: str, dirs: Tuple[str] = (".git",), default=None) -> Optional[str]:
+    """Find the root of the git repository.
+    'https://stackoverflow.com/questions/22081209/find-the-root-of-the-git-repository-where-the-file-lives'
 
     Args:
-        path (_type_): _description_
-        dirs (tuple, optional): _description_. Defaults to (".git",).
-        default (_type_, optional): _description_. Defaults to None.
+        path: The path to start from.
+        dirs: The directories to look for. Defaults to (".git",).
+        default: The default value to return if no git root is found.
 
     Returns:
-        _type_: _description_
+        The root of the git repository. Returns None if not found.
     """
     prev, test = None, os.path.abspath(path)
     while prev != test:
@@ -74,9 +88,12 @@ def get_git_root(path, dirs=(".git",), default=None):
     return default
 
 
-def get_project_root(path):
+def get_project_root(path: str):
     """Return the project root directory from an environment variable.
     # TODO: handle this better and report user error
+
+    Args:
+        path: The path to start from.
     """
     project_root = os.getenv("PROJECT_ROOT")
     if project_root is None:
@@ -103,9 +120,12 @@ def get_absolute_path(path, proj_root_func=get_project_root):
     return absolute_path
 
 
-def make_dir(filename_or_folder):
+def make_dir(filename_or_folder: str) -> str:
     """Make the directory for either the filename or folder.
     Note that filename_or_folder currently needs to end in / for it to be recognized as a folder.
+
+    Args:
+        filename_or_folder (str): The filename or folder to make.
     """
     if filename_or_folder[-1] != "/" and filename_or_folder.find(".") < 0:
         folder = os.path.dirname(filename_or_folder + "/")

pyrad/utils/plotly.py

@@ -42,30 +42,30 @@ def color_str(color):
 
 
 def get_line_segments_from_lines(
-    lines,
-    color=color_str((1, 0, 0)),
-    marker_color=color_str((1, 0, 0)),
-    colors=None,
-    draw_marker=True,
-    draw_line=True,
-    marker_size=4,
-    line_width=10,
-):
+    lines: TensorType["num_rays", 2, 3],
+    color: str = color_str((1, 0, 0)),
+    marker_color: str = color_str((1, 0, 0)),
+    colors: List[str] = None,
+    draw_marker: bool = True,
+    draw_line: bool = True,
+    marker_size: float = 4,
+    line_width: float = 10,
+) -> go.Scatter3d:
     """Returns a list of Scatter3D objects for creating lines with plotly.
     # TODO(ethan): make this function more efficient instead of having a list of objects.
 
     Args:
-        lines (_type_): list of [np.array(2, 3), ...], line is defined by a (2, 3) array
-        color (_type_, optional): _description_. Defaults to color_str((1, 0, 0)).
-        marker_color (_type_, optional): _description_. Defaults to color_str((1, 0, 0)).
-        colors (_type_, optional): _description_. Defaults to None.
-        draw_marker (bool, optional): _description_. Defaults to True.
-        draw_line (bool, optional): _description_. Defaults to True.
-        marker_size (int, optional): _description_. Defaults to 4.
-        line_width (int, optional): _description_. Defaults to 10.
+        lines: Tensor of lines.
+        color: Color of the lines. Defaults to red.
+        marker_color: Color of the markers. Defaults to red.
+        colors: List of colors for each line. Defaults to None.
+        draw_marker: Whether to draw markers. Defaults to True.
+        draw_line: Whether to draw lines. Defaults to True.
+        marker_size: Size of the markers. Defaults to 4.
+        line_width: Width of the lines. Defaults to 10.
 
     Returns:
-        _type_: _description_
+        Scatter3D object on lines.
     """
     data = []
     for idx, line in enumerate(lines):
@@ -92,8 +92,16 @@ def get_line_segments_from_lines(
     return data
 
 
-def visualize_dataset(camera_origins, ray_bundle: RayBundle):
-    """Visualize a dataset with plotly using our cameras and generated rays."""
+def visualize_dataset(camera_origins: TensorType["num_cameras", 3], ray_bundle: RayBundle) -> go.FigureWidget:
+    """Visualize a dataset with plotly using our cameras and generated rays.
+
+    Args:
+        camera_origins: Tensor of camera origins.
+        ray_bundle: Ray bundle.
+
+    Returns:
+        plotly figure.
+    """
 
     skip = 1
     size = 8
@@ -136,11 +144,11 @@ def get_random_color(colormap: List[str] = None, idx: int = None) -> str:
     """Get a random color from a colormap.
 
     Args:
-        colormap (List[str], optional): List of colors. Defaults to Plotly colors.
-        idx (int, optional): Index of color to return. Defaults to None.
+        colormap: List of colors. Defaults to Plotly colors.
+        idx: Index of color to return. Defaults to None.
 
     Returns:
-        str: random color string
+        random color string
     """
     if colormap is None:
         colormap = ex.colors.qualitative.Plotly
@@ -155,14 +163,14 @@ def get_sphere(
     """Returns a sphere object for plotting with plotly.
 
     Args:
-        radius (float): radius of sphere.
-        center (TensorType[3], optional): center of sphere. Defaults to origin.
-        color (str, optional): color of sphere. Defaults to "black".
-        opacity (float, optional): opacity of sphere. Defaults to 1.0.
-        resolution (int, optional): resolution of sphere. Defaults to 32.
+        radius: radius of sphere.
+        center: center of sphere. Defaults to origin.
+        color: color of sphere. Defaults to "black".
+        opacity: opacity of sphere. Defaults to 1.0.
+        resolution: resolution of sphere. Defaults to 32.
 
     Returns:
-        go.Mesh3d: sphere object.
+        sphere object.
     """
     phi = torch.linspace(0, 2 * torch.pi, resolution)
     theta = torch.linspace(-torch.pi / 2, torch.pi / 2, resolution)
@@ -201,16 +209,16 @@ def get_gaussian_ellipsiod(
     """Get a plotly ellipsoid for a Gaussian.
 
     Args:
-        mean (TensorType[3]): mean of the Gaussian.
-        cov: (TesnorType[3]): covariance of the Gaussian.
-        n_std (int, optional): Standard devation to visualize. Defaults to 2 (95% confidence).
-        color (str, optional): Color of the ellipsoid. Defaults to None.
-        opacity (float, optional): Opacity of the ellipsoid. Defaults to 0.5.
-        resolution (int, optional): Resolution of the ellipsoid. Defaults to 20.
-        name (str, optional): Name of the ellipsoid. Defaults to "ellipse".
+        mean: mean of the Gaussian.
+        cov: covariance of the Gaussian.
+        n_std: Standard devation to visualize. Defaults to 2 (95% confidence).
+        color: Color of the ellipsoid. Defaults to None.
+        opacity: Opacity of the ellipsoid. Defaults to 0.5.
+        resolution: Resolution of the ellipsoid. Defaults to 20.
+        name: Name of the ellipsoid. Defaults to "ellipse".
 
     Returns:
-        go.Mesh3d: ellipsoid object.
+        ellipsoid object.
     """
 
     phi = torch.linspace(0, 2 * torch.pi, resolution)
@@ -257,10 +265,10 @@ def get_gaussian_ellipsoids_list(
         gaussians (Gaussians): Gaussians to visualize.
         opacity (float, optional): Opacity of the mesh. Defaults to 0.3.
         color (str, optional): Color of the mesh. Defaults to "random".
-        resolution (int, optional): Resolution of the mesh. Defaults to 20.
+        resolution: Resolution of the mesh. Defaults to 20.
 
     Returns:
-        List[go.Mesh3d]: List of plotly meshes
+        List of plotly meshes
     """
     data = []
 
@@ -297,13 +305,13 @@ def get_frustum_mesh(
     """Get a plotly mesh for a single frustum.
 
     Args:
-        frustum (Frustum): Single frustum
-        opacity (float, optional): Opacity of the mesh. Defaults to 0.3.
-        color (str, optional): Color of the mesh. Defaults to "#DC203C".
-        resolution (int, optional): Resolution of the mesh. Defaults to 20.
+        frustum: Single frustum
+        opacity: Opacity of the mesh. Defaults to 0.3.
+        color: Color of the mesh. Defaults to "#DC203C".
+        resolution: Resolution of the mesh. Defaults to 20.
 
     Returns:
-        go.Mesh3d: Plotly mesh
+        Plotly mesh
     """
 
     if frustum.ndim > 1:
@@ -348,13 +356,13 @@ def get_frustums_mesh_list(
     """Get a list of plotly meshes for a list of frustums.
 
     Args:
-        frustums (Frustum): List of frustums
-        opacity (float, optional): Opacity of the mesh. Defaults to 0.3.
-        color (str, optional): Color of the mesh. Defaults to "random".
-        resolution (int, optional): Resolution of the mesh. Defaults to 20.
+        frustums: List of frustums
+        opacity: Opacity of the mesh. Defaults to 0.3.
+        color: Color of the mesh. Defaults to "random".
+        resolution: Resolution of the mesh. Defaults to 20.
 
     Returns:
-        List[go.Mesh3d]: List of plotly meshes
+        List of plotly meshes
     """
     data = []
     for i, frustum in enumerate(frustums.flatten()):
@@ -372,13 +380,13 @@ def get_frustum_points(
     """Get a set plotly points for frustums centers.
 
     Args:
-        frustum (Frustum): Frustums to visualize.
-        opacity (float, optional): Opacity of the points. Defaults to 0.3.
-        color (str, optional): Color of the poinst. Defaults to "forestgreen".
-        size (float, optional): Size of points. Defaults to 10.
+        frustum: Frustums to visualize.
+        opacity: Opacity of the points. Defaults to 0.3.
+        color: Color of the poinst. Defaults to "forestgreen".
+        size: Size of points. Defaults to 10.
 
     Returns:
-        go.Scatter3d: Plotly points
+        Plotly points
     """
 
     frustum = frustum.flatten()
@@ -404,13 +412,13 @@ def get_ray_bundle_lines(
     """Get a plotly line for a ray bundle.
 
     Args:
-        ray_bundle (RayBundle): Ray bundle
-        length (float, optional): Length of the line. Defaults to 1.0.
-        color (str, optional): Color of the line.
-        width (float, optional): Width of the line. Defaults to 1.
+        ray_bundle: Ray bundle
+        length: Length of the line. Defaults to 1.0.
+        color: Color of the line.
+        width: Width of the line. Defaults to 1.
 
     Returns:
-        go.Scatter3d: Plotly lines
+        Plotly lines
     """
 
     origins = ray_bundle.origins.view(-1, 3)

scripts/run_train.py

@@ -116,9 +116,9 @@ def _train(local_rank: int, world_size: int, config: Config) -> Any:
     """Main training function that sets up and runs the trainer per process
 
     Args:
-        local_rank (int): _description_
         local_rank (int): current rank of process
         world_size (int): total number of gpus available
+        config (Config): config file specifying training regimen
 
     Returns:
         Any: TODO(): determine the return type