content: # Copyright (c) Meta Platforms, Inc. and affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  from typing import Any, Dict, Optional, Type, Union
  
  import numpy as np
  import torch
  from scipy.sparse import csr_matrix
  
  from theseus.core import Objective
  from theseus.optimizer import Linearization, SparseLinearization
  from theseus.optimizer.autograd import BaspachoSolveFunction
  
  from .linear_solver import LinearSolver
  from .utils import convert_to_alpha_beta_damping_tensors
  
  # assuming the want to use cuda if supported and available
  DEFAULT_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
  
  
  class BaspachoSparseSolver(LinearSolver):
      def __init__(
          self,
          objective: Objective,
          linearization_cls: Optional[Type[Linearization]] = None,
          linearization_kwargs: Optional[Dict[str, Any]] = None,
          **kwargs,
      ):
          linearization_cls = linearization_cls or SparseLinearization
          if not linearization_cls == SparseLinearization:
              raise RuntimeError(
                  "BaspachoSparseSolver only works with "
                  "theseus.optimizer.SparseLinearization,"
                  + f" got {type(self.linearization)}"
              )
  
          super().__init__(objective, linearization_cls, linearization_kwargs, **kwargs)
          self.linearization: SparseLinearization = self.linearization
  
          self.device = ""
          self._has_been_reset = False
          if self.linearization.structure().num_rows:
              self.reset()
  
      def _resolve_device_str(self, device: Optional[str] = None) -> str:
          if device is None:
              raw_device = self.linearization.objective.device
              if raw_device is None:
                  device = "cpu"
              elif isinstance(raw_device, torch.device):
                  device = raw_device.type
              else:
                  device = raw_device
          return device
  
      def reset(self, device: Optional[str] = None, **kwargs):
          dev = device or self._resolve_device_str()
          needs_reset = dev != self.device or not self._has_been_reset
          if not needs_reset:
              return
          self.device = dev
          self._has_been_reset = True
          if "cuda" in dev and not torch.cuda.is_available():
              raise RuntimeError(
                  "BaspachoSparseSolver: Cuda requested (dev='cuda') but not\n"
                  "available in torch, use dev='cpu' to create a Cpu-based solver"
              )
  
          try:
              from theseus.extlib.baspacho_solver import SymbolicDecomposition
          except Exception as e:
              raise RuntimeError(
                  "Theseus C++ extension cannot be loaded.\n"
                  "This is an installation issue (also check CUDA\n"
                  "if CUDA support was compiled in)"
                  f"{type(e).__name__}: {e}"
              )
  
          # convert to tensors for accelerated Mt x M operation
          self.A_row_ptr = torch.tensor(
              self.linearization.structure().row_ptr, dtype=torch.int64
          ).to(dev)
          self.A_col_ind = torch.tensor(
              self.linearization.structure().col_ind, dtype=torch.int64
          ).to(dev)
  
          # compute block-structure of AtA. To do so we multiply the Jacobian's
          # transpose At by a matrix that collapses the rows to block-rows, ie
          # if i-th block starts at b_i, then a non-zero in b_i-th row becomes
          # a non-zero in i-th row.
          At_mock = self.linearization.structure().mock_csc_transpose()
          num_vars = len(self.linearization.var_start_cols)
          num_cols = self.linearization.num_cols
          to_blocks = csr_matrix(
              (
                  np.ones(num_cols),
                  np.arange(num_cols),
                  [*self.linearization.var_start_cols, num_cols],
              ),
              (num_vars, num_cols),
          )
          block_At_mock = to_blocks @ At_mock
          block_AtA_mock = (block_At_mock @ block_At_mock.T).tocsr()
          block_AtA_mock.sort_indices()
  
          param_size = torch.tensor(self.linearization.var_dims, dtype=torch.int64)
          block_struct_ptrs = torch.tensor(block_AtA_mock.indptr, dtype=torch.int64)
          block_struct_inds = torch.tensor(block_AtA_mock.indices, dtype=torch.int64)
          self.symbolic_decomposition = SymbolicDecomposition(
              param_size, block_struct_ptrs, block_struct_inds, dev
          )
  
      def solve(
          self,
          damping: Optional[Union[float, torch.Tensor]] = None,
          ellipsoidal_damping: bool = True,
          damping_eps: float = 1e-8,
          **kwargs,
      ) -> torch.Tensor:
          if not isinstance(self.linearization, SparseLinearization):
              raise RuntimeError(
                  "CholmodSparseSolver only works with theseus.optimizer.SparseLinearization."
              )
  
          if damping is None:
              damping_alpha_beta = None
          else:
              damping_alpha_beta = convert_to_alpha_beta_damping_tensors(
                  damping,
                  damping_eps,
                  ellipsoidal_damping,
                  batch_size=self.linearization.A_val.shape[0],
                  device=self.linearization.A_val.device,
                  dtype=self.linearization.A_val.dtype,
              )
  
          return BaspachoSolveFunction.apply(
              self.linearization.A_val,
              self.linearization.b,
              self.linearization.structure(),
              self.A_row_ptr,
              self.A_col_ind,
              self.symbolic_decomposition,
              damping_alpha_beta,
              self.linearization.detached_hessian,
          )
  