I have a convex optimization problem


\begin{align\*}
\max\_{P\, =\, (p\_{ij})} &\sum\_{ij} p\_{ij} \big(a\_{ij}-\log(p\_{ij})\big)\\
 \text{st}\quad &p\_{ij}\in\mathcal{P}\_{n, m} \; \text{and} \sum\_{i=1}^n p\_{ij} (b\_i-1) = 0 \quad \forall j\\
 \text{where,}\; &\mathcal{P}\_{n, m} := \{(p\_{ij}) \in \mathbb{R}\_+^{n \times m}:\; \sum\_{j} p\_{ij} = 1,\forall i\} 
\end{align\*}


This problem is related to [A convex optimization problem](https://mathoverflow.net/questions/255982/a-convex-optimization-problem). For that problem it was suggested to use proximal operator and it worked well but for this one the proximal map doesn't have a closed-form. How to go about this?


A bit of search shows that there are these generalized proximal mappings where one replaces squared Euclidean distance with other distance-like functions. The keywords here seem to be "entropic proximal mapping" and "mirror descent". Are there versions of splitting algorithms such as Douglas-Rachford for these generalized proximal operators?



Proximal map for a convex optimization problem
(This is my second attempt at a solution. My initial attempt can be found below.)


Your optimization problem can be expressed as
$$
\operatorname{minimize}\_{p \in \mathbb R^{n \times m}} \quad F(p) + G(p)
$$
where
$$
F(p) = \sum\_{i,j} p\_{ij}(\log(p\_{ij}) - a\_{ij}) + I\_\Omega(p)
$$
and $I\_\Omega$ is the indicator function of the set
$$
\Omega = \{ p \in \mathbb R^{n \times m} \mid p\_{ij} \geq 0 \, \forall \, i,j\}
$$
(so $I\_\Omega$ enforces nonnegativity on $p$)
and
$G$ is the indicator function of the set
$$
S = \{ p \in \mathbb R^{n \times m} \mid
\sum\_j p\_{ij} = 1 \, \forall \, i \quad \text{and}\quad
\sum\_{i=1}^n p\_{ij}(b\_i - 1) = 0 \, \forall j \}.
$$
Evaluating the prox-operator of $G$ requires projecting onto $S$, which is a linear algebra problem with a standard closed-form solution.


We do not have a closed-form expression for the prox-operator of $F$, **but** $F$ is still a very simple function (for example, it is fully separable). Can we not evaluate the prox-operator of $F$ numerically to high precision very efficiently?


Evaluating the prox-operator of $F$ reduces to evaluating the prox-operator of the function $w:\mathbb R \to \mathbb R \cup \{ \infty \}$ defined by
$$
w(x) = \begin{cases} x \log(x) & \quad \text{if } x \geq 0, \\
\infty & \quad \text{otherwise.}
\end{cases}
$$
(We interpret $x \log(x) = 0$ when $x = 0$.)
To evaluate the prox-operator of $w$ at $\hat x$ (with parameter $t > 0$), we must solve
\begin{align}
\operatorname{minimize}\_x & \quad x \log(x) + \frac{1}{2t} (x - \hat x)^2 \\
\text{subject to} & \quad x \geq 0.
\end{align}
If we visualize the graph of $ v(x) = x \log(x) + \frac{1}{2t} (x - \hat x)^2$, we see that $v$ is initially decreasing (as we move away from the origin, to the right) but then eventually increases to $+\infty$. So, $v$ has a minimizer in the interval $(0,\infty)$ which can be found by setting $v'(x) = 0$. This yields:
$$
1 + \log(x) + \frac{1}{t}(x - \hat x) = 0
\iff x + t \log(x) = \hat x - t.
$$
The function $x + t \log(x)$ is strictly increasing, and it ranges from $-\infty$ to $+\infty$ as $x$ ranges from $0$ to $+\infty$.
Thus, there is a single value of $x$ for which
$x + t \log(x) = \hat x - t$.
I think this value of $x$ can be found to high accuracy with Newton's method (and I think it won't take too many iterations). This allows us to evaluate the prox-operator of $F$ efficiently.


Alternatively (and even better), we can express the solution to 
$$
\frac{x}{t} + \log(x) = c
$$
in terms of the [Lambert W-function](https://en.wikipedia.org/wiki/Lambert_W_function).
Note that
\begin{align}
\frac{x}{t} + \log(x) = c & \iff
\log(e^{x/t}) + \log(x) = \log(e^c) \\
&\iff x e^{x/t} = e^c \\
&\iff (x/t) e^{x/t} = \frac{e^c}{t}\\
&\iff x/t = W(e^c/t) \\
&\iff x = t W(e^c/t).
\end{align}
In Matlab, the Lambert W-function can be evaluated using `lambertw`.


Now that we have seen how to evaluate the prox-operators of both $F$ and $G$ efficiently, we are able to minimize $F(p) + G(p)$ using the Douglas-Rachford method (which solves convex problems of exactly this form).




---


Here is my initial attempt to solve the problem, which had a certain flaw (or issue, at least) which is discussed at the end.


The Chambolle-Pock algorithm solves convex optimization problems in the canonical form (or "graph form") minimize $f(x) + g(Cx)$, where the convex functions $f$ and $g$ are "simple" (meaning they have easy prox-operators) and $C$ is a matrix or a linear transformation. The Chambolle-Pock algorithm was extended by Condat (in [this paper](http://link.springer.com/article/10.1007/s10957-012-0245-9)) to solve convex optimization problems of the form
$$
\tag{$\spadesuit$} \operatorname{minimize}\_x \quad h(x) + f(x) + g(Cx)
$$
where $f,g$ and $C$ are as above and the convex function $h$ is differentiable with a Lipschitz continuous gradient. Your problem has the form $(\spadesuit)$ where $C$ is the identity mapping and
\begin{align}
h(p) = -\sum\_{ij} \, p\_{ij}(a\_{ij} - \log(p\_{ij})) \\
\end{align}
and $f$ is the indicator function of
$$
\{ p \in \mathbb R^{n \times m}\_+ \mid \sum\_j p\_{ij} = 1 \, \forall i \}
$$
and
$g$ is the indicator function of
$$
\{ p \in \mathbb R^{m \times n} \mid \sum\_{i=1}^n p\_{ij}(b\_i - 1) = 0 \, \forall j \}.
$$


**NOTE:** A **problem** (or at least an issue) with this approach is that $h$ does not have a Lipschitz continuous gradient and cannot be extended to a differentiable function $\mathbb R^{n \times m}$. That might mean this approach doesn't work.


I'd be interested to see what @MichaelGrant says about this problem.


