SOLAX: A Python solver for fermionic quantum systems with neural network support

1- The codebase addresses a clear need with well-justified design decision and modular implementation.
2- The code is based on existing packages and frameworks and integrates them well. The installation procedure seems quite simple.
3- The code is compatible "out-of-the-box" with existing NVIDIA GPUs easily enabling parallelization.
4- The provided documentation is clear and easy to follow, the detailed examples are appreciated.

1- I find that the paper could be more self-contained by not relying so much on some of the cited works in which methods were introduced and providing more context.
2- The paper would benefit from more discussion of limitations and future directions/updates.

This paper introduces the SOLAX Python package which provides utilities for working with fermionic quantum states and operators in the second quantization. It allows to construct and manipulate arbitrary basis sets, and express quantum states and operators in these bases. It leverages existing standard libraries such as Numpy, SciPy and JAX and is compatible with NVIDIA GPUs for parallelization. It also includes utilities to select basis states in large Hilbert spaces using neural networks and thus improve the efficiency of ground-state computations.

I think this codebase is a nice contribution and will likely be useful amongst practitioners, especially those interested in applications of deep learning to ground-state computations. Being a reader familiar with the machine learning side more than the condensed matter one, I found that the examples given could benefit from a bit more discussion on the physics (see requests). I would also appreciate some more on potential limitations and potential problem cases of the machine-learning basis selection method.

The article would be well-suited for publication in SciPost Physics Codebases. I would recommend its acceptance there with minor modifications. I don't think it meets the criteria for SciPost Physics since it does not present a new method or results.

1- Are there libraries that implement similar functionalities? If yes, what are the main differences in capabilities and design choices between the proposed libraries and alternative ones?
2- I would like to see a more detailed discussion of what the main future directions would be so that potential users have an idea of what to look for in coming releases or things they could contribute.
3- For operators, explain quickly dagger convention as there are multiple possible orderings.
4- If no batch size is specified, will the GPU be used by default? This should be clarified.
5 - "We assume that this cutoff also divides the full set of candidates in the same proportion." This seems like a strong assumption. Is it really necessary to make it? If yes, can the authors provide an argument for why this assumption is reasonable. A trivial failure case is for example when then ground state is one of the basis states. Are there other likely ones?
6- To make the paper more self-explanatory can some more motivation for the deep learning approach can be provided? Why should it be expected to work well and generalize beyond the training set?
7- What are potential limitations and potential problem cases of the machine-learning basis selection method? There should be a discussion of these.
8- I did not understand why a random instance is generated each time a random sample is needed. Some more details on this in 3.2.1 would be good.
9- Can the authors provide more discussion about the relevance of the SIAM problem? It would also be necessary to provide some short justification on the choice of Hamiltonian parameters.
10- Can the authors provide some more discussion on the basis extension method along with references? What are some strengths and limitations of this method?

The following are not requested changes, but features to consider:
- It would be nice to have a "fancy" printing option for operators to display them in a more readable format (sum of creation-annihilation operators products).
- Adding the possibility to label the orbitals in some meaningful way could be an interesting feature, to allow the neural network to rely on these labels. For example, explicitly what is the spin value, or angular momentum quantum number associated with a given orbital. Leveraging this information instead of just the position of the bits in the Slater determinant could help the neural network make better predictions.

Accept in alternative Journal (see Report)

1 - Easy to install
2 - In particular the solver is well documented and its usage is straightforward
3 - The userguide is easy to read and provides examples to get started
4 - The library makes use of JAX to allow for Neural Network (NN) support

1 - The NN support is providing very little background and uses machine learning jargon which needs further explanation
2 - In a similar vain, the example calculation (solving the Single Impurity Anderson Model, SIAM) needs more context

This manuscript provides an introduction to SOLAX, a python library for solving fermionic quantum systems (in particular impurity models) with neural network support. In my opinion, this manuscript and the corresponding code are well suited for a publication in SciPost Physics Codebases after some corrections and changes as detailed below. In particular, I think that the acceptance criteria for a Codebases article are met if further benchmarking tests are provided with the source code. However, I do not think that it should be published as a SciPost Physics article, the category under which it has been (perhaps by mistake) submitted.
The corresponding program code can be installed rather easily, but the authors might still want to add a few remarks concerning the required versions of the solax and orbax packages (see ‘requested changes’ below).

The userguide is divided into three main parts: A description of the solver, the neural network support, and the saving and loading of SOLAX objects.
The largest part of the userguide, chapter 2, focusses on the functionalities of the solver, in particular on the Basis, State, OperatorTerm, Operator, and OperatorMatrix classes. The different functionalities of these classes are nicely explained in a pedagogical manner that does not require specialized knowledge and is very well suited for users to get started with the SOLAX library. Ample example code snippets illustrate the functionalities and jupyter notebooks are provided for the three main chapters of the text.
Part three is, however, hard to follow through for a non-expert. Especially readers with a background in solid state physics or quantum chemistry without prior detailed knowledge of machine learning and deep learning concepts might experience difficulties. This is unfortunate since it could be easily the target readership of this article. I think the paper would benefit largely from a revision which aims at better explaining the key concepts and providing enough information to understand this chapter without having to read the NN books mentioned in the introduction or the authors’ paper on solving the SIAM, Ref. [16]. I understand that certain basic concepts using NN can be expected from a user of SOLAX, but a more pedagogical approach in chapter 3 as detailed in the requested changes below would still be desirable.

Concerning the pros and cons of the presented solver, it is clearly stated in the conclusions that the main advantage of SOLAX is to allow treating large bath sizes of up to 299 bath sites in the context of the SIAM. However, the authors should also comment on the limitations of the solver and mention how much computation time is spent on the training of the NN in the two mentioned cases (SIAM, N$_2$ molecule) . Furthermore it would be useful for a reader to know in which situations a large number of sites is actually needed for the presented (single-orbital) SIAM.

1 - In the installation section, the authors specify which versions of the required python libraries have been used to run the example scripts shown in the userguide. This is perfectly fine. However, it might be useful to point out a few delicate points concerning the flax and orbax libraries. In particular when using the Anaconda ecosystem, the provided library versions are not sufficient to run the jupyter notebooks correctly. Also, it might be worth pointing out to opt for the most recent orbax-checkpoint version.
2 - In section 2.2.3 A, please specify how the chop method handles complex coefficients of the State objects.
3 - As a suggestion to improve the readability of equations (1), (2), etc., adding the corresponding formulas in human readable notation might help, i.e. $a^{\dagger}_{1,\uparrow}$ instead of $a_{0}^{\dagger}$.
In the same spirit, I would suggest to add some more physical context whenever it is possible, e.g. on page 16, section 2.3.4, where a singlet state is encoded.
4 - On page 17 experimental features are mentioned. Please say explicitly which features are tested and which are experimental in the current version of the code!
In the same spirit, it is mentioned on page 15, section 2.3.1 that there are other “manual mode” features which are not detailed in the userguide. Where are these functions documented?
5 - The example matrix on pages 23/24ff. seems to have a wrong diagonal part. At least when following the notation in the userguide (and also the jupyter notebook) I obtain for the specified basis a different potential energy due to the U term: on the diagonal (1.0, 1.25, 1.0, 1.75). The example matrix should be corrected.
6 - Concerning the code snippets in the userguide, on page 5, sec. 2.1, solax and numpy are imported. It would be consistent to import there already scipy, which is needed later for the sparse eigensolver on page 34.
7 - In Figure 1 the level scheme is slightly confusing. From this figure it is not clear that the hybridizations $V_b$ and energies $\epsilon_b$ are different for the different bath sites. In particular the “energy arrow” needs to be updated in order to avoid the impression that all bath levels are situated at the same energy.
8 - Page 29: Please add a brief motivation and introduce the density of states (semi-circular) and the hybridization (cosine) used in the example calculation. A more compact version of the intro in Sec. III of Ref. 16 around eq. (16) might be a good starting point.
Page 30: Similarly, briefly motivate and explain the choice of GS via two Slater determinants (compare to the passage around eq. (17) of Ref.16).
9 - Concerning chapter 3, I think in particular the function nn_call_on_bits needs better explanation (section 3.2, page 39f.). Please avoid abbreviations like ReLU (rectified linear unit), explain the dense block, comment on why there is no padding in the convolutional layer, etc. Also, a figure which explains the architecture of the NN such as Fig. 2 of Ref. 16 might help.
10 - In chapter 3, it would be good to clearly separate parameters with formal justification from those which have been optimized, fine-tuned or chosen based on “machine learning experience”. In particular the choices of the following parameters are not clear:
(a) dense_size – explain the generic expression used here ;
(b) It is not clear on which basis the number of most important determinants has been determined (target_num on page 41). This needs to be better explained.
(c) Why taking 2/3 of the target set size for the random set used for training? Is this factor of 2/3 fine-tuned or can it be justified in a formal way?
11 - Two minor corrections:
- page 36, center “[…] passing both Basis objects[…]” (without 'the')
- page 43, first sentence: “[...]for conveniently storing[...]”

Accept in alternative Journal (see Report)