1: #include <petsc/private/taoimpl.h>

  3: /*@
  4:   TaoSoftThreshold - Calculates soft thresholding routine with input vector
  5:   and given lower and upper bound and returns it to output vector.

  7:   Collective

  9:   Input Parameters:
 10: + in - input vector to be thresholded
 11: . lb - lower bound
 12: - ub - upper bound

 14:   Output Parameter:
 15: . out - Soft thresholded output vector

 17:   Notes:
 18:   Soft thresholding is defined as
 19:   \[ S(input,lb,ub) =
 20:   \begin{cases}
 21:   input - ub  & \text{if } input > ub \\
 22:   0           & \text{if } lb \leq input \leq ub \\
 23:   input - lb  & \text{if } input < lb
 24:   \end{cases}
 25:   \]

 27:   Level: developer

 29: .seealso: `Tao`, `Vec`
 30: @*/
 31: PetscErrorCode TaoSoftThreshold(Vec in, PetscReal lb, PetscReal ub, Vec out)
 32: {
 33:   PetscInt     i, nlocal, mlocal;
 34:   PetscScalar *inarray, *outarray;

 36:   PetscFunctionBegin;
 39:   PetscCheck(lb <= ub, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Lower bound needs to be lower than upper bound");
 40:   if (lb == ub) {
 41:     PetscCall(VecCopy(in, out));
 42:     PetscCall(VecShift(out, -lb));
 43:     PetscFunctionReturn(PETSC_SUCCESS);
 44:   }

 46:   PetscCall(VecGetLocalSize(in, &nlocal));
 47:   PetscCall(VecGetLocalSize(out, &mlocal));
 48:   PetscCheck(nlocal == mlocal, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Input and output vectors need to be of same size");
 49:   PetscCall(VecGetArrayPair(in, out, &inarray, &outarray));

 51:   for (i = 0; i < nlocal; i++) outarray[i] = PetscMax(0, PetscRealPart(inarray[i]) - ub) + PetscMin(0, PetscRealPart(inarray[i]) - lb);

 53:   PetscCall(VecRestoreArrayPair(in, out, &inarray, &outarray));
 54:   PetscFunctionReturn(PETSC_SUCCESS);
 55: }