text: # Copyright 2021 The NetKet Authors - All rights reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #    http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  from typing import Union, Optional, Tuple, Any, Callable, Iterable
  
  import numpy as np
  
  import jax
  from jax import numpy as jnp
  from flax import linen as nn
  
  from netket.hilbert import AbstractHilbert
  from netket.graph import AbstractGraph
  from netket.utils.types import PRNGKey, Shape, Dtype, Array, NNInitFunc
  
  from netket import nn as nknn
  from netket.nn.initializers import lecun_normal, variance_scaling, zeros, normal
  
  from .rbm import RBM
  
  default_kernel_init = normal(stddev=0.001)
  
  
  class PureRBM(nn.Module):
      """
      Encodes the pure state |ψ><ψ|, because it acts on row and column
      indices with the same RBM.
      """
  
      dtype: Any = np.float64
      """The dtype of the weights."""
      activation: Any = nknn.logcosh
      """The nonlinear activation function."""
      alpha: Union[float, int] = 1
      """feature density. Number of features equal to alpha * input.shape[-1]"""
      use_hidden_bias: bool = True
      """if True uses a bias in the dense layer (hidden layer bias)."""
      use_visible_bias: bool = True
      """if True adds a bias to the input not passed through the nonlinear layer."""
      precision: Any = None
      """numerical precision of the computation see `jax.lax.Precision`for details."""
  
      kernel_init: NNInitFunc = default_kernel_init
      """Initializer for the Dense layer matrix."""
      hidden_bias_init: NNInitFunc = zeros
      """Initializer for the hidden bias."""
      visible_bias_init: NNInitFunc = zeros
      """Initializer for the visible bias."""
  
      @nn.compact
      def __call__(self, σr, σc, symmetric=True):
          W = nknn.Dense(
              name="Dense",
              features=int(self.alpha * σr.shape[-1]),
              dtype=self.dtype,
              use_bias=self.use_hidden_bias,
              kernel_init=self.kernel_init,
              bias_init=self.hidden_bias_init,
              precision=self.precision,
          )
          xr = self.activation(W(σr)).sum(axis=-1)
          xc = self.activation(W(σc)).sum(axis=-1)
  
          if symmetric:
              y = xr + xc
          else:
              y = xr - xc
  
          if self.use_visible_bias:
              v_bias = self.param(
                  "visible_bias", self.visible_bias_init, (σr.shape[-1],), self.dtype
              )
              if symmetric:
                  out_bias = jnp.dot(σr + σc, v_bias)
              else:
                  out_bias = jnp.dot(σr - σc, v_bias)
  
              y = y + out_bias
  
          return 0.5 * y
  
  
  class MixedRBM(nn.Module):
      """
      Encodes the pure state |ψ><ψ|, because it acts on row and column
      indices with the same RBM.
      """
  
      dtype: Any = np.float64
      """The dtype of the weights."""
      activation: Any = nknn.logcosh
      """The nonlinear activation function."""
      alpha: Union[float, int] = 1
      """feature density. Number of features equal to alpha * input.shape[-1]"""
      use_bias: bool = True
      """if True uses a bias in the dense layer (hidden layer bias)."""
      precision: Any = None
      """numerical precision of the computation see `jax.lax.Precision`for details."""
  
      kernel_init: NNInitFunc = default_kernel_init
      """Initializer for the Dense layer matrix."""
      bias_init: NNInitFunc = zeros
      """Initializer for the hidden bias."""
  
      @nn.compact
      def __call__(self, σr, σc, symmetric=True):
          U_S = nknn.Dense(
              name="Symm",
              features=int(self.alpha * σr.shape[-1]),
              dtype=self.dtype,
              use_bias=False,
              kernel_init=self.kernel_init,
              precision=self.precision,
          )
          U_A = nknn.Dense(
              name="ASymm",
              features=int(self.alpha * σr.shape[-1]),
              dtype=self.dtype,
              use_bias=False,
              kernel_init=self.kernel_init,
              precision=self.precision,
          )
          y = U_S(0.5 * (σr + σc)) + 1j * U_A(0.5 * (σr - σc))
  
          if self.use_bias:
              bias = self.param(
                  "bias",
                  self.bias_init,
                  (int(self.alpha * σr.shape[-1]),),
                  jax.dtypes.dtype_real(self.dtype),
              )
              y = y + bias
  
          y = self.activation(y)
          return y.sum(axis=-1)
  
  
  class NDM(nn.Module):
      """
      Encodes a Positive-Definite Neural Density Matrix using the ansatz from Torlai and
      Melko, PRL 120, 240503 (2018).
  
      Assumes real dtype.
      A discussion on the effect of the feature density for the pure and mixed part is
      given in Vicentini et Al, PRL 122, 250503 (2019).
      """
  
      dtype: Any = np.float64
      """The dtype of the weights."""
      activation: Any = nknn.logcosh
      """The nonlinear activation function."""
      alpha: Union[float, int] = 1
      """The feature density for the pure-part of the ansatz. 
      Number of features equal to alpha * input.shape[-1]
      """
      beta: Union[float, int] = 1
      """The feature density for the mixed-part of the ansatz. 
      Number of features equal to beta * input.shape[-1]
      """
      use_hidden_bias: bool = True
      """if True uses a bias in the dense layer (hidden layer bias)."""
      use_ancilla_bias: bool = True
      """if True uses a bias in the dense layer (hidden layer bias)."""
      use_visible_bias: bool = True
      """if True adds a bias to the input not passed through the nonlinear layer."""
      precision: Any = None
      """numerical precision of the computation see `jax.lax.Precision`for details."""
  
      kernel_init: NNInitFunc = default_kernel_init
      """Initializer for the Dense layer matrix."""
      bias_init: NNInitFunc = default_kernel_init
      """Initializer for the hidden bias."""
      visible_bias_init: NNInitFunc = default_kernel_init
      """Initializer for the visible bias."""
  
      @nn.compact
      def __call__(self, σ):
          σr, σc = jax.numpy.split(σ, 2, axis=-1)
  
          ψ_S = PureRBM(
              name="PureSymm",
              alpha=self.alpha,
              activation=self.activation,
              dtype=self.dtype,
              use_hidden_bias=self.use_hidden_bias,
              use_visible_bias=self.use_visible_bias,
              visible_bias_init=self.visible_bias_init,
              kernel_init=self.kernel_init,
              hidden_bias_init=self.bias_init,
              precision=self.precision,
          )
  
          ψ_A = PureRBM(
              name="PureASymm",
              alpha=self.alpha,
              activation=self.activation,
              dtype=self.dtype,
              use_hidden_bias=self.use_hidden_bias,
              use_visible_bias=self.use_visible_bias,
              visible_bias_init=self.visible_bias_init,
              kernel_init=self.kernel_init,
              hidden_bias_init=self.bias_init,
              precision=self.precision,
          )
  
          Π = MixedRBM(
              name="Mixed",
              alpha=self.beta,
              dtype=self.dtype,
              use_bias=self.use_ancilla_bias,
              activation=self.activation,
              kernel_init=self.kernel_init,
              bias_init=self.bias_init,
              precision=self.precision,
          )
          return (
              ψ_S(σr, σc, symmetric=True) + 1j * ψ_A(σr, σc, symmetric=False) + Π(σr, σc)
          )
  
          x = self.activation(x)
  
          return jnp.sum(x, axis=-1) + out_bias
  