[WIP] Adds specializations for handling sparse matrix unary operations

stan/math/mix.hpp

@@ -5,15 +5,15 @@
 #include <stan/math/mix/fun.hpp>
 #include <stan/math/mix/functor.hpp>
 
-#ifdef STAN_OPENCL
-#include <stan/math/opencl/rev.hpp>
-#endif
-
 #include <stan/math/fwd/core.hpp>
 #include <stan/math/fwd/meta.hpp>
 #include <stan/math/fwd/fun.hpp>
 #include <stan/math/fwd/functor.hpp>
 
+#ifdef STAN_OPENCL
+#include <stan/math/opencl/rev.hpp>
+#endif
+
 #include <stan/math/rev/core.hpp>
 #include <stan/math/rev/meta.hpp>
 #include <stan/math/rev/fun.hpp>

stan/math/prim/fun/acos.hpp

@@ -47,7 +47,7 @@ template <typename Container,
           require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
               Container>* = nullptr>
 inline auto acos(const Container& x) {
-  return apply_scalar_unary<acos_fun, Container>::apply(x);
+  return apply_scalar_unary<acos_fun, Container, true>::apply(x);
 }
 
 /**

stan/math/prim/fun/cos.hpp

@@ -42,7 +42,7 @@ template <typename Container,
           require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
               Container>* = nullptr>
 inline auto cos(const Container& x) {
-  return apply_scalar_unary<cos_fun, Container>::apply(x);
+  return apply_scalar_unary<cos_fun, Container, true>::apply(x);
 }
 
 /**

stan/math/prim/fun/digamma.hpp

@@ -75,7 +75,7 @@ template <typename T,
           require_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr,
           require_not_var_matrix_t<T>* = nullptr>
 inline auto digamma(const T& x) {
-  return apply_scalar_unary<digamma_fun, T>::apply(x);
+  return apply_scalar_unary<digamma_fun, T, true>::apply(x);
 }
 
 }  // namespace math

stan/math/prim/fun/erfc.hpp

@@ -37,7 +37,7 @@ template <
     require_all_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr,
     require_not_var_matrix_t<T>* = nullptr>
 inline auto erfc(const T& x) {
-  return apply_scalar_unary<erfc_fun, T>::apply(x);
+  return apply_scalar_unary<erfc_fun, T, true>::apply(x);
 }
 
 }  // namespace math

stan/math/prim/fun/exp.hpp

@@ -45,7 +45,7 @@ template <
     require_not_nonscalar_prim_or_rev_kernel_expression_t<Container>* = nullptr,
     require_not_var_matrix_t<Container>* = nullptr>
 inline auto exp(const Container& x) {
-  return apply_scalar_unary<exp_fun, Container>::apply(x);
+  return apply_scalar_unary<exp_fun, Container, true>::apply(x);
 }
 
 /**

stan/math/prim/fun/exp2.hpp

@@ -40,7 +40,7 @@ template <
     require_all_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr,
     require_not_var_matrix_t<T>* = nullptr>
 inline auto exp2(const T& x) {
-  return apply_scalar_unary<exp2_fun, T>::apply(x);
+  return apply_scalar_unary<exp2_fun, T, true>::apply(x);
 }
 
 }  // namespace math

stan/math/prim/fun/inv.hpp

@@ -35,7 +35,7 @@ template <
     typename T, require_not_container_st<std::is_arithmetic, T>* = nullptr,
     require_all_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr>
 inline auto inv(const T& x) {
-  return apply_scalar_unary<inv_fun, T>::apply(x);
+  return apply_scalar_unary<inv_fun, T, true>::apply(x);
 }
 
 /**

stan/math/prim/fun/inv_Phi.hpp

@@ -178,7 +178,7 @@ template <
     require_all_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr,
     require_not_var_matrix_t<T>* = nullptr>
 inline auto inv_Phi(const T& x) {
-  return apply_scalar_unary<inv_Phi_fun, T>::apply(x);
+  return apply_scalar_unary<inv_Phi_fun, T, true>::apply(x);
 }
 
 }  // namespace math

stan/math/prim/fun/inv_cloglog.hpp

@@ -77,7 +77,7 @@ template <typename Container,
           require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
               Container>* = nullptr>
 inline auto inv_cloglog(const Container& x) {
-  return apply_scalar_unary<inv_cloglog_fun, Container>::apply(x);
+  return apply_scalar_unary<inv_cloglog_fun, Container, true>::apply(x);
 }
 
 /**

stan/math/prim/fun/inv_logit.hpp

@@ -85,7 +85,7 @@ template <
     typename T, require_not_var_matrix_t<T>* = nullptr,
     require_all_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr>
 inline auto inv_logit(const T& x) {
-  return apply_scalar_unary<inv_logit_fun, T>::apply(x);
+  return apply_scalar_unary<inv_logit_fun, T, true>::apply(x);
 }
 
 // TODO(Tadej): Eigen is introducing their implementation logistic() of this

stan/math/prim/fun/inv_sqrt.hpp

@@ -46,7 +46,7 @@ template <typename Container,
           require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
               Container>* = nullptr>
 inline auto inv_sqrt(const Container& x) {
-  return apply_scalar_unary<inv_sqrt_fun, Container>::apply(x);
+  return apply_scalar_unary<inv_sqrt_fun, Container, true>::apply(x);
 }
 
 /**

stan/math/prim/fun/lgamma.hpp

@@ -117,7 +117,7 @@ struct lgamma_fun {
 template <typename T, require_not_var_matrix_t<T>* = nullptr,
           require_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr>
 inline auto lgamma(const T& x) {
-  return apply_scalar_unary<lgamma_fun, T>::apply(x);
+  return apply_scalar_unary<lgamma_fun, T, true>::apply(x);
 }
 
 }  // namespace math

stan/math/prim/fun/log.hpp

@@ -48,7 +48,7 @@ template <
     require_not_var_matrix_t<Container>* = nullptr,
     require_not_nonscalar_prim_or_rev_kernel_expression_t<Container>* = nullptr>
 inline auto log(const Container& x) {
-  return apply_scalar_unary<log_fun, Container>::apply(x);
+  return apply_scalar_unary<log_fun, Container, true>::apply(x);
 }
 
 /**

stan/math/prim/fun/log10.hpp

@@ -39,7 +39,7 @@ template <
     require_not_container_st<std::is_arithmetic, Container>* = nullptr,
     require_not_nonscalar_prim_or_rev_kernel_expression_t<Container>* = nullptr>
 inline auto log10(const Container& x) {
-  return apply_scalar_unary<log10_fun, Container>::apply(x);
+  return apply_scalar_unary<log10_fun, Container, true>::apply(x);
 }
 
 /**

stan/math/prim/fun/log1m.hpp

@@ -71,7 +71,7 @@ template <
     typename T, require_not_var_matrix_t<T>* = nullptr,
     require_all_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr>
 inline auto log1m(const T& x) {
-  return apply_scalar_unary<log1m_fun, T>::apply(x);
+  return apply_scalar_unary<log1m_fun, T, true>::apply(x);
 }
 
 }  // namespace math

stan/math/prim/fun/log1m_exp.hpp

@@ -81,7 +81,7 @@ template <
     typename T, require_not_var_matrix_t<T>* = nullptr,
     require_all_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr>
 inline auto log1m_exp(const T& x) {
-  return apply_scalar_unary<log1m_exp_fun, T>::apply(x);
+  return apply_scalar_unary<log1m_exp_fun, T, true>::apply(x);
 }
 
 }  // namespace math

stan/math/prim/fun/log1m_inv_logit.hpp

@@ -82,9 +82,8 @@ struct log1m_inv_logit_fun {
  */
 template <typename T, require_not_var_matrix_t<T>* = nullptr,
           require_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr>
-inline typename apply_scalar_unary<log1m_inv_logit_fun, T>::return_t
-log1m_inv_logit(const T& x) {
-  return apply_scalar_unary<log1m_inv_logit_fun, T>::apply(x);
+inline auto log1m_inv_logit(const T& x) {
+  return apply_scalar_unary<log1m_inv_logit_fun, T, true>::apply(x);
 }
 
 }  // namespace math

stan/math/prim/fun/log1p.hpp

@@ -80,7 +80,7 @@ template <typename T,
           require_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr,
           require_not_var_matrix_t<T>* = nullptr>
 inline auto log1p(const T& x) {
-  return apply_scalar_unary<log1p_fun, T>::apply(x);
+  return apply_scalar_unary<log1p_fun, T, true>::apply(x);
 }
 
 }  // namespace math

stan/math/prim/fun/log1p_exp.hpp

@@ -76,7 +76,7 @@ template <typename T,
           require_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr,
           require_not_var_matrix_t<T>* = nullptr>
 inline auto log1p_exp(const T& x) {
-  return apply_scalar_unary<log1p_exp_fun, T>::apply(x);
+  return apply_scalar_unary<log1p_exp_fun, T, true>::apply(x);
 }
 
 }  // namespace math

stan/math/prim/fun/log2.hpp

@@ -47,7 +47,7 @@ struct log2_fun {
 template <typename T, require_not_var_matrix_t<T>* = nullptr,
           require_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr>
 inline auto log2(const T& x) {
-  return apply_scalar_unary<log2_fun, T>::apply(x);
+  return apply_scalar_unary<log2_fun, T, true>::apply(x);
 }
 
 }  // namespace math

stan/math/prim/fun/log_inv_logit.hpp

@@ -81,7 +81,7 @@ struct log_inv_logit_fun {
 template <typename T,
           require_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr>
 inline auto log_inv_logit(const T& x) {
-  return apply_scalar_unary<log_inv_logit_fun, T>::apply(x);
+  return apply_scalar_unary<log_inv_logit_fun, T, true>::apply(x);
 }
 
 }  // namespace math

stan/math/prim/fun/logit.hpp

@@ -89,7 +89,7 @@ template <
     require_not_var_matrix_t<Container>* = nullptr,
     require_not_nonscalar_prim_or_rev_kernel_expression_t<Container>* = nullptr>
 inline auto logit(const Container& x) {
-  return apply_scalar_unary<logit_fun, Container>::apply(x);
+  return apply_scalar_unary<logit_fun, Container, true>::apply(x);
 }
 
 /**

stan/math/prim/fun/tgamma.hpp

@@ -51,7 +51,7 @@ template <
     require_all_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr,
     require_not_var_matrix_t<T>* = nullptr>
 inline auto tgamma(const T& x) {
-  return apply_scalar_unary<tgamma_fun, T>::apply(x);
+  return apply_scalar_unary<tgamma_fun, T, true>::apply(x);
 }
 
 }  // namespace math

stan/math/prim/functor/apply_scalar_unary.hpp

@@ -2,12 +2,7 @@
 #define STAN_MATH_PRIM_FUNCTOR_APPLY_SCALAR_UNARY_HPP
 
 #include <stan/math/prim/fun/Eigen.hpp>
-#include <stan/math/prim/meta/is_eigen.hpp>
-#include <stan/math/prim/meta/is_complex.hpp>
-#include <stan/math/prim/meta/require_generics.hpp>
-#include <stan/math/prim/meta/is_vector.hpp>
-#include <stan/math/prim/meta/is_vector_like.hpp>
-#include <stan/math/prim/meta/plain_type.hpp>
+#include <stan/math/prim/meta.hpp>
 #include <utility>
 #include <vector>
 
@@ -36,8 +31,12 @@ namespace math {
  *
  * @tparam F Type of function to apply.
  * @tparam T Type of argument to which function is applied.
+ * @tparam ApplyZero If true, the function applied is assumed to return zero for
+ *  inputs of zero and so sparse matrices will return sparse matrices. A value
+ * of false will return a dense matrix for sparse matrices.
  */
-template <typename F, typename T, typename Enable = void>
+template <typename F, typename T, bool ApplyZero = false,
+          typename Enable = void>
 struct apply_scalar_unary;
 
 /**
@@ -48,8 +47,8 @@ struct apply_scalar_unary;
  * @tparam F Type of function to apply.
  * @tparam T Type of argument to which function is applied.
  */
-template <typename F, typename T>
-struct apply_scalar_unary<F, T, require_eigen_t<T>> {
+template <typename F, typename T, bool ApplyZero>
+struct apply_scalar_unary<F, T, ApplyZero, require_eigen_t<T>> {
   /**
    * Type of underlying scalar for the matrix type T.
    */
@@ -59,15 +58,82 @@ struct apply_scalar_unary<F, T, require_eigen_t<T>> {
    * Return the result of applying the function defined by the
    * template parameter F to the specified matrix argument.
    *
+   * @tparam DenseMat A type derived from `Eigen::DenseBase`.
    * @param x Matrix to which operation is applied.
    * @return Componentwise application of the function specified
    * by F to the specified matrix.
    */
-  static inline auto apply(const T& x) {
+  template <typename DenseMat, require_eigen_dense_base_t<DenseMat>* = nullptr>
+  static inline auto apply(const DenseMat& x) {
     return x.unaryExpr(
         [](scalar_t x) { return apply_scalar_unary<F, scalar_t>::apply(x); });
   }
 
+  /**
+   * Special case for `ApplyZero` set to true, returning a full sparse matrix.
+   * Return the result of applying the function defined by the template
+   * parameter F to the specified matrix argument.
+   *
+   * @param SparseMat A type derived from `Eigen::SparseMatrixBase`
+   * @tparam NonZeroZero Shortcut trick for using class template for deduction,
+   * should not be set manually.
+   * @param x Matrix to which operation is applied.
+   * @return Componentwise application of the function specified
+   * by F to the specified matrix.
+   */
+  template <typename SparseMat, bool NonZeroZero = ApplyZero,
+            require_t<bool_constant<NonZeroZero>>* = nullptr,
+            require_eigen_sparse_base_t<SparseMat>* = nullptr>
+  static inline auto apply(const SparseMat& x) {
+    using val_t = value_type_t<SparseMat>;
+    using triplet_t = Eigen::Triplet<val_t>;
+    auto zeroed_val = apply_scalar_unary<F, scalar_t>::apply(val_t(0.0));
+    const auto x_size = x.size();
+    std::vector<triplet_t> triplet_list(x_size, triplet_t(0, 0, zeroed_val));
+    for (Eigen::Index i = 0; i < x.rows(); ++i) {
+      for (Eigen::Index j = 0; j < x.cols(); ++j) {
+        // Column major order
+        triplet_list[i * x.cols() + j] = triplet_t(i, j, zeroed_val);
+      }
+    }
+    for (Eigen::Index k = 0; k < x.outerSize(); ++k) {
+      for (typename SparseMat::InnerIterator it(x, k); it; ++it) {
+        triplet_list[it.row() * x.cols() + it.col()]
+            = triplet_t(it.row(), it.col(),
+                        apply_scalar_unary<F, scalar_t>::apply(it.value()));
+      }
+    }
+    plain_type_t<SparseMat> ret(x.rows(), x.cols());
+    ret.setFromTriplets(triplet_list.begin(), triplet_list.end());
+    return ret;
+  }
+
+  /**
+   * Special case for `ApplyZero` set to false, returning a sparse matrix.
+   * Return the result of applying the function defined by the template
+   * parameter F to the specified matrix argument.
+   *
+   * @tparam SparseMat A type derived from `Eigen::SparseMatrixBase`
+   * @tparam NonZeroZero Shortcut trick for using class template for deduction,
+   * should not be set manually.
+   * @param x Matrix to which operation is applied.
+   * @return Componentwise application of the function specified
+   * by F to the specified matrix.
+   */
+  template <typename SparseMat, bool ReturnZeros = ApplyZero,
+            require_t<bool_constant<!ReturnZeros>>* = nullptr,
+            require_eigen_sparse_base_t<SparseMat>* = nullptr>
+  static inline auto apply(const SparseMat& x) {
+    auto ret = x.eval();
+    for (Eigen::Index k = 0; k < x.outerSize(); ++k) {
+      for (typename SparseMat::InnerIterator it(x, k), ret_it(ret, k); it;
+           ++it, ++ret_it) {
+        ret_it.valueRef() = apply_scalar_unary<F, scalar_t>::apply(it.value());
+      }
+    }
+    return ret;
+  }
+
   /**
    * Return type for applying the function elementwise to a matrix
    * expression template of type T.
@@ -82,8 +148,8 @@ struct apply_scalar_unary<F, T, require_eigen_t<T>> {
  *
  * @tparam F Type of function defining static apply function.
  */
-template <typename F, typename T>
-struct apply_scalar_unary<F, T, require_floating_point_t<T>> {
+template <typename F, typename T, bool ApplyZero>
+struct apply_scalar_unary<F, T, ApplyZero, require_floating_point_t<T>> {
   /**
    * The return type, double.
    */
@@ -107,8 +173,8 @@ struct apply_scalar_unary<F, T, require_floating_point_t<T>> {
  *
  * @tparam F Type of function defining static apply function.
  */
-template <typename F, typename T>
-struct apply_scalar_unary<F, T, require_complex_t<T>> {
+template <typename F, typename T, bool ApplyZero>
+struct apply_scalar_unary<F, T, ApplyZero, require_complex_t<T>> {
   /**
    * The return type, double.
    */
@@ -134,8 +200,8 @@ struct apply_scalar_unary<F, T, require_complex_t<T>> {
  *
  * @tparam F Type of function defining static apply function.
  */
-template <typename F, typename T>
-struct apply_scalar_unary<F, T, require_integral_t<T>> {
+template <typename F, typename T, bool ApplyZero>
+struct apply_scalar_unary<F, T, ApplyZero, require_integral_t<T>> {
   /**
    * The return type, double.
    */
@@ -162,8 +228,8 @@ struct apply_scalar_unary<F, T, require_integral_t<T>> {
  * @tparam F Class defining a static apply function.
  * @tparam T Type of element contained in standard vector.
  */
-template <typename F, typename T>
-struct apply_scalar_unary<F, std::vector<T>> {
+template <typename F, typename T, bool ApplyZero>
+struct apply_scalar_unary<F, std::vector<T>, ApplyZero, void> {
   /**
    * Return type, which is calculated recursively as a standard
    * vector of the return type of the contained type T.