Extension of PBC code to 1D about ferminet HOT 7 CLOSED

from typing import Optional, Tuple

import chex
from ferminet import networks
import jax.numpy as jnp

def periodic_norm_1d(metric: jnp.ndarray, scaled_r: jnp.ndarray) -> jnp.ndarray:
    del metric #unused
    return jnp.sin(jnp.pi*scaled_r[..., 0])/jnp.pi#Divide by pi to recover the euclidean norm in the limit of small scaled_r


#Same function as in the original code, just uses my norm.
def make_pbc_feature_layer_1d(
    natoms: Optional[int] = None,
    nspins: Optional[Tuple[int, ...]] = None,
    ndim: int = 1,
    rescale_inputs: bool = False,
    lattice: Optional[jnp.ndarray] = None,
    include_r_ae: bool = False,
) -> networks.FeatureLayer:
  """Returns the init and apply functions for periodic features.

  Args:
      natoms: number of atoms.
      nspins: tuple of the number of spin-up and spin-down electrons.
      ndim: dimension of the system.
      rescale_inputs: If true, rescales r_ae for stability. Note that unlike in
        the OBC case, we do not rescale r_ee as well.
      lattice: Matrix whose columns are the primitive lattice vectors of the
        system, shape (ndim, ndim).
      include_r_ae: Flag to enable electron-atom distance features. Set to False
        to avoid cusps with ghost atoms in, e.g., homogeneous electron gas.
  """

  del nspins

  if lattice is None:
    lattice = jnp.eye(ndim)

  # Calculate reciprocal vectors, factor 2pi omitted
  reciprocal_vecs = jnp.linalg.inv(lattice)
  lattice_metric = lattice.T @ lattice

  def init() -> Tuple[Tuple[int, int], networks.Param]:
    if include_r_ae:
      return (natoms * (2 * ndim + 1), 2 * ndim + 1), {}
    else:
      return (natoms * (2 * ndim), 2 * ndim + 1), {}

  def apply(ae, r_ae, ee, r_ee) -> Tuple[jnp.ndarray, jnp.ndarray]:
    # One e features in phase coordinates, (s_ae)_i = k_i . ae
    s_ae = jnp.einsum('il,jkl->jki', reciprocal_vecs, ae)
    # Two e features in phase coordinates
    s_ee = jnp.einsum('il,jkl->jki', reciprocal_vecs, ee)
    # Periodized features
    ae = jnp.concatenate(
        (jnp.sin(2 * jnp.pi * s_ae), jnp.cos(2 * jnp.pi * s_ae)), axis=-1)
    ee = jnp.concatenate(
        (jnp.sin(2 * jnp.pi * s_ee), jnp.cos(2 * jnp.pi * s_ee)), axis=-1)
    # Distance features defined on orthonormal projections
    r_ae = periodic_norm_1d(lattice_metric, s_ae)
    if rescale_inputs:
      r_ae = jnp.log(1 + r_ae)
    # Don't take gradients through |0|
    n = ee.shape[0]
    s_ee += jnp.eye(n)[..., None]
    r_ee = periodic_norm_1d(lattice_metric, s_ee) * (1.0 - jnp.eye(n))

    if include_r_ae:
      ae_features = jnp.concatenate((r_ae[..., None], ae), axis=2)
    else:
      ae_features = ae
    ae_features = jnp.reshape(ae_features, [jnp.shape(ae_features)[0], -1])
    ee_features = jnp.concatenate((r_ee[..., None], ee), axis=2)
    return ae_features, ee_features

  return networks.FeatureLayer(init=init, apply=apply)

from typing import Mapping, Optional, Sequence, Tuple, Union
import jax.numpy as jnp

def make_kpoints_1d(
    lattice: Union[jnp.ndarray, jnp.ndarray],
    spins = Tuple[int, int],
    min_kpoints: Optional[int] = None
) -> jnp.ndarray:
    if min_kpoints is None:
        min_kpoints = sum(spins)
    elif min_kpoints < sum(spins):
        raise ValueError("Number of kpoints must be greater than number of electrons")
    
    #kpts will be of the form i*mod, with i an integer between -n and n inclusive
    mod = 2*jnp.pi/lattice[0][0]
    if min_kpoints%2 == 1: n = min_kpoints
    else: n = min_kpoints + 1 #In 1d always generate an odd number of pts
    n //= 2

    kpoints = mod*jnp.arange(-n, n+1, 1)
    #If the first kpt is zero the envelope evaluates to 1 at the beginning of training
    kpoints = jnp.concatenate((kpoints[n:], kpoints[:n])) 
    return kpoints[:, None]

from ferminet import networks
from ferminet import hamiltonian
import jax.numpy as jnp
from typing import Sequence

def local_energy_free_particles(
    f: networks.FermiNetLike,
    charges: jnp.ndarray,
    nspins: Sequence[int],
    use_scan: bool = False,
    complex_output: bool = False,
) -> hamiltonian.LocalEnergy:
    """Just a wrapper with a signature that matches the one of 
    hamiltonian.MakeLocalEnergy"""

    del nspins, charges
    ke = hamiltonian.local_kinetic_energy(
        f,
        use_scan=use_scan,
        complex_output=complex_output
    )
    def _e_l(params, key, data):
        del key
        kinetic = ke(params, data)
        return kinetic

    return _e_l

import sys
import subprocess as sp
import os
import time

from absl import logging
import jax
import jax.numpy as jnp

from ferminet import base_config as cfg
from ferminet.utils import system as system
from ferminet import train
from ferminet.pbc import envelopes

from pbc_1d import envelopes_1d as env
from pbc_1d import feature_layer_1d as fl
from pbc_1d import hamiltonian_1d as ham

def main():
    # Optional, for also printing training progress to STDOUT.
    # If running a script, you can also just use the --alsologtostderr flag.
    logging.get_absl_handler().python_handler.stream = sys.stdout
    logging.set_verbosity(logging.INFO)

    lattice = jnp.array([[5]]) #Matrix with lattice vectors as columns
    placeholder_atom = system.Atom(
      symbol = "X", #"Element" with atomic number 0
      charge = 0,
      coords = (0.,)
    )

    my_config = cfg.default()

    my_config.system.ndim = 1
    my_config.system.molecule = [placeholder_atom]
    systems = [(1, 1)]

    my_config.pretrain.iterations = 0
    my_config.optim.lr.rate = 1e-4

    my_config.network.make_feature_layer_fn = "pbc_1d.feature_layer_1d.make_pbc_feature_layer_1d"
    my_config.network.make_feature_layer_kwargs = {"include_r_ae":False, "lattice":lattice}
    my_config.network.make_envelope_fn = "ferminet.pbc.envelopes.make_multiwave_envelope"
    my_config.system.make_local_energy_fn = "pbc_1d.hamiltonian_1d.local_energy_free_particles"
    #The following is only used for the potential energy
    #my_config.system.make_local_energy_kwargs = {"lattice": lattice}
    my_config.network.determinants = 16
    my_config.log.save_frequency = 5.
    my_config.optim.iterations = 100

    my_config.debug.deterministic = True

    A = 1e4
    C = 1.
    B = 1
    
    my_config.optim.lr.rate = B/(A**C)
    my_config.optim.lr.decay = C
    my_config.optim.lr.delay = A

    for spin_pair in systems:
        print("Running spin_pair", spin_pair)
        my_config.system.electrons = spin_pair
        #If even, make_kpoints will add one
        min_kpoints = spin_pair[0] + spin_pair[1] + 1
        kpoints = env.make_kpoints_1d(lattice, my_config.system.electrons, min_kpoints = min_kpoints)
        print("kpoints:", kpoints)
        my_config.network.make_envelope_kwargs = {"kpoints":kpoints}
        train.train(my_config)

if __name__ == "__main__":
    main()