design question for a custom operator

[dambuck]

Hi,

I need to compute a observable which was problematic to implement using LocalOperator because of the sheer number of operators I need #457 . There i was told that a custom operator like nk.operator.Ising would be a good choice. The observables I need have the form
Sz_i H Sz_j, where H is a 1D Heisenberg hamiltonian with 2 sites per uni cell. Before I tried to build one LocalOp for every (i,j) which simply took too much memory. So after going through the custom operator docs I just wanted to assure myself that I am choosing a sensible way to implement the custom observables.

the way I understand it, I will need to define at least the following: get_conns_and_mels(), e_loc(), get_local_kernel_arguments() and get_local_kernel(). For Sz_i H Sz_j I need 2*n_sites observables, so I would like to define them all within one custom Op. This would mean that in e.g. the tutorial function

def get_conns_and_mels(sigma):
  # this code only works if sigma is a single bitstring
  assert sigma.ndim == 1

  # get number of spins
  N = sigma.shape[-1]
  # repeat eta N times
  eta = jnp.tile(sigma, (N,1))
  # diagonal indices
  ids = np.diag_indices(N)
  # flip those indices
  eta = eta.at[ids].set(-eta.at[ids].get())
  return eta, jnp.ones(N)

that returns connected states and the corresponding matrix elements for get_local_kernel() I need to add another axis to account for the different observables (i,j).I wanted to follow "Defining an operator the easy way" but am I right that the default expect as in netket/vqs/mc/mc_state/expect.py would fail or can I somehow @jax.vmap around it? I am also slightly lost because of the many dispatch rules for expect and the different kinds of nk.stats.statistics /mpi.statistics calls that appear and which to use in a custom _expect.
I guess the way to go is to set up a new dispatch rule to account for another axis in the arguments for get_local_kernel()?

I would be very thankful for some tipps on this

[PhilipVinc]

OperatorLean as presented in the custom operator documentation basically ask you to provide a function through get_local_kernel that evaluates the local energy for a batch of configurations, and to prepare it's input with get_local_kernel_argument.

If you have different observables (I,j) you can just flatten the output you return from get_local_kernel_argument?

In short, I think you should be able to take the example in here and edit only get_local_kernel_arguments to return the connected elements of all entries.

[PhilipVinc]

You cannot add another axis, because the Lean interface is built with this limitation. But you can flatten it yourself.

[PhilipVinc]

An alternative would be to implement an operator inheriting from DiscreteOperator, like Ising:

netket/netket/operator/_hamiltonian.py

Line 252 in 216092e

def _flattened_kernel(

In this case you'd have to define a function (ideally numbajitted) that computes the connected elements.
As those operators correspond to a matrix, again, you cannot add new axes arbitrarily but you can flatten it yourself.

The point of not adding a new dimension is because, in general, an operator represent the matrix (X_1HX_1 + X_1 H X_2 + X_2H X_1 + X_2 H X_2 + etc...) , and not the single terms.

[PhilipVinc]

Are you at the Toulouse school by any chance? It's hard to follow what you want to do.
The only thing that I can say is that if you want to do something different from what stats does (for example, to have a stats object that correctly propagates errors) you need to rewrite the whole expect function (aka, you can't use OperatorLean, which relies on returning a stats object).

You can build your own 'stats' object, and use it in your own expect dispatch for a custom operator. as long as it also has a .mean field it will work with everything else well.

[dambuck]

thanks for the quick reply again. Sadly I could not make it to Toulouse.
I explained poorly but what I want is essentially, instead of doing: vs.expect(operator) with vs.n_samples=100 that returns the stats object; I would like to accesss the 100 local values of O(x) myself. Whatever error propagation I need, ill do locally on my machine with those samples.

[femtobit]

@dambuck Well, your timing for that is pretty good, because returning the local energies from .expect is something we were working on right now. See #1154 for a draft.

[singhavishek]

@dambuck Can I have a look at the code that you wrote for custom operator. I am also writing a similar code and having some problems.

[dambuck]

Hi,

its been a while since I used it, but here is the releant part. Its pretty ugly but it worked :p

class SHS(DiscreteOperator):
    def __init__(
        self,
        hilbert: AbstractHilbert,
        graph: AbstractGraph,
        sign_rule =None,
        J: float = 1.0,
        J_: float=.5,
        pairing: int =0,
        distance: int =0):
        
        super().__init__(hilbert)

        self._edges = np.asarray(
            [[u, v] for u, v in graph.edges()],
            dtype=np.intp,)
        
        
            
        
        self.J=float(J)
        self.J_=float(J_)
        self.pairing=pairing
        self.distance=distance
        
        
        assert distance<graph.n_nodes//2 , 'invalid distance for lattice'
        assert pairing in [0,1,2,3], 'invalid pairing'
        
        if sign_rule is None:
                self.sign_rule = [1] * hilbert.size
        else:
            assert len(sign_rule) == hilbert.size
            self.sign_rule=sign_rule
              
        

    @property
    def dtype(self):
        return complex

    @staticmethod
    @jit(nopython=True)
    def _flattened_kernel(
        x,
        sections,
        edges,
        sign_rule,
        J,
        J_,
        pair,
        distance,
        pad=True):
        
        
        n_sites = x.shape[1]
        # max conn are in case of a sample [-1,+1,-1,+1,...]
        #n_sites//2 for S+S- , n_sites//2 for S-S+, N_sites//4 for NNN S+S-
        # n_sites//4 for NNN S-S+
        # +1 for the diagonal part which just preserves the sampled state
        max_conn=n_sites+n_sites//2 +1 
        conn_batch=np.zeros(x.shape[0],dtype=np.intp) #number of conn states for every batch
        NN_pm_conn=np.zeros(x.shape[0],dtype=np.intp) #number of conn states for S+S- op
        NN_mp_conn=np.zeros(x.shape[0],dtype=np.intp) #number of conn states for S-S+ op
        NNN_mp_conn=np.zeros(x.shape[0],dtype=np.intp)
        NNN_pm_conn=np.zeros(x.shape[0],dtype=np.intp)
        # number of opposite spins, NN and NNN
        tot_conn=0
        for batch in range(x.shape[0]):
            conn_i=0
            for i in range(x.shape[1]):
                if x[batch,i]*x[batch,(i+1)%n_sites]<0:
                    conn_i+=1
                    if x[batch,i]>0:
                        NN_pm_conn[batch]+=1
                    else:
                        NN_mp_conn[batch]+=1
                if i%2==0:
                    if x[batch,i]*x[batch,(i+2)%n_sites] <0:
                        conn_i+=1
                        if x[batch,i]>0:
                            NNN_pm_conn[batch]+=1
                        else:
                            NNN_mp_conn[batch]+=1
                            
            assert (NN_pm_conn[batch]+NN_mp_conn[batch]\
                    + NNN_pm_conn[batch]+NNN_mp_conn[batch])== conn_i, 'tot conn not equal summed up conn'
            conn_i+=1  #for the diagonal part of NN and NNN
            
            tot_conn+=conn_i
            
            
            # !!!!!  PADDING OR NOT SECTIONS
            # sections[batch]=tot_conn
            conn_batch[batch]=conn_i
            
        
        
        
        # number of NN edges
        # n_NN_edges= edges.shape[0]//3 *2
    
        x_prime = np.zeros((tot_conn,n_sites))
        mels = np.zeros( tot_conn, dtype=type(J)) #! only works with type(), else numba error
#        mels = np.zeros( tot_conn, dtype=sign_rule.dtype) #! only works with type(), else numba error

        if pad:
            for batch in range(x.shape[0]):
                sections[batch]=max_conn
            
            x_prime = np.zeros((x.shape[0]*max_conn,n_sites))
            mels = np.zeros( x.shape[0]*max_conn, dtype=type(J))

        site_shift=[0,0,-1,1]
        loop_range=[range(0,x.shape[1],2),range(1,x.shape[1],2),range(1,x.shape[1],2),range(0,x.shape[1],2)]
        

        diag_ind = 0
        
        
        for batch in range(x.shape[0]):
            # diagonal connected state
            x_prime[diag_ind]=np.copy(x[batch])
            # print(mels)
            # print('diag in', diag_ind)
            ###### compute DIAGONAL matrix elements #########
            for trans in loop_range[pair]:
                S_k=trans
                S_l=(trans+2*distance +site_shift[pair])%n_sites
            # S_k=0
            # S_l=0
                ##### looping through contributions in the Hamiltonian #####
                NN_edge_counter=0
                _NN_edge_counter=0
                NNN_edge_counter=0
                _NNN_edge_counter=0
                for start in range(edges.shape[0]):
                    # sign phase, carefull at boundary edges!
                    # check if boundary tuple or not
                    if abs(edges[start,1]-edges[start,0])<3:
                        sign_phase=sign_rule[edges[start,1]]/sign_rule[edges[start,0]] 
                        _sign_phase= sign_rule[edges[start,0]]/sign_rule[edges[start,1]] 
                        ij_reversed=False
                        i=edges[start,0]
                        j=edges[start,1]
                    else:
                        sign_phase=sign_rule[edges[start,0]]/sign_rule[edges[start,1]]
                        _sign_phase= sign_rule[edges[start,1]]/sign_rule[edges[start,0]] 
                        ij_reversed=True
                        i=edges[start,1]
                        j=edges[start,0]
                        
                        
                        # REMEMBER TO CHANGE
                    # NN=False
                    
                    if i%2!=0 or j%2!=0:
                        # print('NN',i, j)
                        # print(i%2 ,j%2)
                        NN=True
                    else:
                        # print('NNN',i, j)
                        # print(i%2 ,j%2)
                        NN=False
                    
                    # spins parallel, no contrib from +- operators
                    if x[batch,i]* x[batch,j]>0:
                        # Sz_k  Sz_i Sz_j   Sz_l DIAG part for NN AND NNN
                        if NN:
                            mels[diag_ind]+= J *x[batch,S_k]* x[batch,S_l] * x[batch,i] * x[batch,j]*0.25 * 0.25* 1/(n_sites//2)
                            
                        else:
                            mels[diag_ind]+= J_\
                                        *x[batch,S_k]* x[batch,S_l] * x[batch,i] * x[batch,j]*0.25 * 0.25 *1/(n_sites//2)
                                       
                        # print(x[batch], diag_ind,mels[diag_ind], (i,j))
                        continue
                    
                    
                    
                    # NN diag part, NN edges are bewtween NN sites
                    if NN:
                        # print('i,j', (i,j))
                        # Sz_k  Sz_i Sz_j   Sz_l  DIAGONAL
                        mels[diag_ind]+= J *x[batch,S_k]* x[batch,S_l] * x[batch,i] * x[batch,j]*0.25 * 0.25* 1/(n_sites//2)
                        
               
                        # print('diag part',(i,j) ,J *x[batch,S_k]* x[batch,S_l] * x[batch,i] * x[batch,j]*0.25 * 0.25)
                            
                        #########    S+_i S-_j elements , sigma_i is spin up #############
                        if x[batch,i]>0:
                            # print(i,j)
                            # Sz_k Sz_l S+_i S-_j  , all operators act on differents sites
                            # here the boundary problem with ordered tuples does not matter
                            if S_l!=i and S_l!=j:
                                mels[diag_ind+1 + NN_edge_counter]+=  J/2* x[batch,S_k] *x[batch,S_l]\
                                    * sign_phase*   0.25* 1/(n_sites//2)
                                
                            
                                
                          
                            else:
                                if S_l==j:
                                    mels[diag_ind+1 + NN_edge_counter]+= J/2  * x[batch,S_k]/2 *(x[batch,S_l]/2+1) \
                                        *sign_phase* 1/(n_sites//2)
                                    
                                elif S_l==i:
                                    mels[diag_ind+1 + NN_edge_counter]+= J/2 * x[batch,S_k]/2* (x[batch,S_l]/2-1) \
                                        *sign_phase* 1/(n_sites//2)
                                    
                                        
                            x_prime[diag_ind+1 + NN_edge_counter] = np.copy(x[batch])
                            x_prime[diag_ind+1 + NN_edge_counter, i] *= -1.0
                            x_prime[diag_ind+1 + NN_edge_counter, j] *= -1.0
                            
                            NN_edge_counter+=1
                        #########  END  S+_i S-_j elements #############
                        
                        #########    S-_i S+_j elements, sigma_i is spin down #############
                        else:
                            # Sz_k Sz_l S-_i S+_j  
                            if S_l!=i and S_l!=j:
                                mels[diag_ind+ n_sites//2 +1+ _NN_edge_counter]+= J/2 * x[batch,S_k] *x[batch,S_l]\
                                    * _sign_phase*  0.25* 1/(n_sites//2)
                                    
                               
                         
                            else:
                                if S_l==j:
                                    mels[diag_ind+ n_sites//2 +1+ _NN_edge_counter]+= J/2 * x[batch,S_k]/2*(x[batch,S_l]/2-1) \
                                      *_sign_phase* 1/(n_sites//2)
                                      
                                   
                                elif S_l==i:
                                    mels[diag_ind+ n_sites//2 +1+ _NN_edge_counter]+= J/2 * x[batch,S_k]/2 *(x[batch,S_l]/2+1) \
                                        * _sign_phase* 1/(n_sites//2)
                                        
                                  
                                    
                            x_prime[diag_ind+ n_sites//2 +1+ _NN_edge_counter] = np.copy(x[batch])
                            x_prime[diag_ind+ n_sites//2 +1+ _NN_edge_counter, i] *= -1.0
                            x_prime[diag_ind+ n_sites//2 +1+ _NN_edge_counter, j] *= -1.0
                            _NN_edge_counter+=1
                            
                            # CHECK IF +- OPERATORS BREAK DOWN THE STATE 
                    # NNN contributions
                    else:
                        # print('i,j', (i,j))
                          # Sz_k  Sz_i Sz_j   Sz_l  DIAGONAL
                        mels[diag_ind]+= J_  *x[batch,S_k] * x[batch,S_l] *0.25 \
                                        * x[batch,i] * x[batch,j]*0.25 \
                                        * 1/(n_sites//2)
                        
                        #########    S+_i S-_j elements , sigma_i is spin up #############
                        if x[batch,i]>0:
                            # Sz_k Sz_l S+_i S-_j  , all operators act on differents sites
                            if S_l!=i and S_l!=j:
                                mels[diag_ind + n_sites +1+ NNN_edge_counter ]+=  J_/2 * x[batch,S_k] *x[batch,S_l]\
                                    * sign_phase* 0.25 * 1/(n_sites//2)
                                    
                            # Sz_l_index== S-_j_index  : commute Sz_l to the left
                          
                           
                            else:
                                if S_l==j:
                                    mels[diag_ind +n_sites+1+  NNN_edge_counter]+= J_/2 * x[batch,S_k]/2* (x[batch,S_l]/2+1)\
                                        *sign_phase* 1/(n_sites//2)
                                elif S_l==i:
                                    mels[diag_ind+n_sites+1+  NNN_edge_counter]+= J_/2 * x[batch,S_k]/2* (x[batch,S_l]/2-1)\
                                        *sign_phase* 1/(n_sites//2)
                                        
                            x_prime[diag_ind +n_sites+1+  NNN_edge_counter] = np.copy(x[batch])
                            x_prime[diag_ind +n_sites+1+  NNN_edge_counter, i] *= -1.0
                            x_prime[diag_ind +n_sites+1+  NNN_edge_counter, j] *= -1.0
                            NNN_edge_counter+=1
                            #########    S-_i S+_j elements, sigma_i is spin down #############
                        else:
                            if S_l!=i and S_l!=j:
                                mels[diag_ind + n_sites + n_sites//4 +1 + _NNN_edge_counter ]+=  J_/2* x[batch,S_k] *x[batch,S_l]\
                                      *_sign_phase* 0.25* 1/(n_sites//2)
                            
                            else:
                                if S_l==j:
                                    mels[diag_ind + n_sites + n_sites//4 +1 + _NNN_edge_counter ]+=  J_/2* \
                                        x[batch,S_k]/2* (x[batch,S_l]/2-1)\
                                        *_sign_phase* 1/(n_sites//2)
                                elif S_l==i:
                                    mels[diag_ind + n_sites + n_sites//4 +1 + _NNN_edge_counter ]+=  J_/2* \
                                        x[batch,S_k]/2* (x[batch,S_l]/2+1)\
                                        *_sign_phase* 1/(n_sites//2)
                            
                            x_prime[diag_ind+ n_sites + n_sites//4 +1 + _NNN_edge_counter ] = np.copy(x[batch])
                            x_prime[diag_ind + n_sites + n_sites//4 +1 + _NNN_edge_counter , i] *= -1.0
                            x_prime[diag_ind + n_sites + n_sites//4 +1 + _NNN_edge_counter , j] *= -1.0
                            _NNN_edge_counter+=1
                    
            # diag_ind+=conn_batch[batch]
            # print('batch and prime',x_prime[diag_ind], x[batch])
            diag_ind+=max_conn
            
        

        return x_prime, mels


    def get_conn_flattened(
        self,
        x,
        sections,
        pad=False,
    ):
        r"""Finds the connected elements of the Operator. Starting
        from a given quantum number x, it finds all other quantum numbers x' such
        that the matrix element :math:`O(x,x')` is different from zero. In general there
        will be several different connected states x' satisfying this
        condition, and they are denoted here :math:`x'(k)`, for :math:`k=0,1...N_{\mathrm{connected}}`.
        This is a batched version, where x is a matrix of shape (batch_size,hilbert.size).
        Args:
            x (matrix): A matrix of shape (batch_size,hilbert.size) containing
                        the batch of quantum numbers x.
            sections (array): An array of size (batch_size) useful to unflatten
                        the output of this function.
                        See numpy.split for the meaning of sections.
            pad (bool): no effect here
        Returns:
            matrix: The connected states x', flattened together in a single matrix.
            array: An array containing the matrix elements :math:`O(x,x')` associated to each x'.
        """

        return self._flattened_kernel(
            np.asarray(x),
            sections,
            self._edges,
            self.sign_rule,
            self.J,
            self.J_,
            self.pairing,
            self.distance
        )

[dambuck]

thank you for the quick answer as always!
So as my observable S_iHS_j includes the hamiltonian, I will prob need the graph and the sign structure to compute the matrix elements of the connected states. Could there be potential problems when continuing from the OperatorLean in the docs by redefining an __init__ that takes the graph object as arg?

I guess this makes inheriting DiscreteOperator the cleaner choice? Also is there a reason why the get_conn defined in DiscreteOperator seems to include the state x itself as first element of the array of connected states x' with matrix element 0?

[PhilipVinc]

. Could there be potential problems when continuing from the OperatorLean in the docs by redefining an init that takes the graph object as arg?

not at all. We do exactly this for the Ising hamiltonian.

I guess this makes inheriting DiscreteOperator the cleaner choice?

They are different choices, both valid.
Inheriting from discrete operator forces you to write the get_conn_flattened using Numba, which is often easier but sometimes slower than writing it with Jax, which you can only do in the other cases.
Though in many cases this won't be a bottleneck.

fined in DiscreteOperator seems to include the state x itself as first element of the array of connected states x' with matrix element 0?

Where does this happen?

[dambuck]

it doesnt actually, I made a mistake when I was checking the method outputs