The authors present a method to make a surrogate model of a photochemistry model of low to intermediate complexity, including an interpretable physical constraint in the neural network. The work is very relevant both from the mathematical (machine learning) concept and for the application in atmospheric modelling. Mathematical concepts are combined with physical/chemical interpretations and the demonstration model is well chosen to be relevant by itself and being complex enough to demonstrate the generalizability of the method to larger models.  Methods and technical results are mostly presented in a concise and clear way.

I have some minor comments and suggestions to improve the readability of the paper. The method is demonstrated for a photochemical model which can be an aim by itself. This combination of  the concept with a clear demonstration is on the one hand the strength of the paper. On the other hand, making clear reference to the demonstration case in the earlier (more conceptual) part of the paper can help the reader.

Detailed comment:

Abstract

The context can be narrowed on the one hand (atmospheric composition/air quality) and widened on the other (impact of emission changes AND climate change) on atmospheric composition/air quality. Would be good to more specifically mention what the demonstration model is. Photochemistry is important part of air quality models, can be aim by itself, and in addition this model contains all essential elements for generalization.

Introduction

The authors could point out what is different with atmospheric reactions with resperct to other ML approaches with conservation laws at the beginning of the introduction. (l 55-60) would make a nice part of the introductoin. As a reference to the general context of machine learning in earth science, the paper by Kashinath gives a good perspective  and deserves citation.

Kashinath  et al 2021 Physics-informed machine learning: case studies for weather and climate modelling Phil. Trans. Royal Sociaety 20200093https://doi.org/10.1098/rsta.2020.0093

L37: Reference to Beucler et al 2019 would also  be relevant, where also the neural network architecture itself is used for constraint, not just part of the output cost function. Beucler T, Rasp S, Pritchard M, Gentine P. 2019 Achieving conservation of energy in neura network emulators for climate modeling. (http://arxiv.org/abs/1906.06622).

2 Derivation and model configuration

Figure 1 caption explicity mentions 11 species and additional parameters, but the photochemical model is not mentioned in the text yet. Would also good to mention the demonstration model very briefly in the introduction. Or leave out the explicit reference to the number 11 and the meteorological parameters.

2.3 Would be better to first describe the gas-phase photochemistry model, then the Julia motivation.

O2 is missing from Table 2. Diatiomic oxygen is present at much higher concentrations that 0.209 ppm, makes nearly 21% of troposphere so I suspect at unit error (absolute mixing ratio). Cf. line 329.

Figure 2 caption: Sentence can be misunderstood, what do you mean by ‘data the NNs were not optimized to predict’

Around line 300, Figure 3, I’m confused with  the 23 hours. Is Figure3 equal to Figure 2 but then withouth the first hour of simulation of each simulation day?

3.4: Figure 6 is not present and the text l349 seems to refer to Figure 5.

The authors consider neural network surrogate models with physics constraints to enforce conservation laws, they demonstrate the effectiveness of their method using a photochemistry model. Experiment results show perfect atom conservation. The presentation is clear. It is an important research topic to enforce conservation laws (e.g., mass conservation) in neural network surrogates.

Some comments on the paper:

The authors should refer to and discuss the paper: Beucler et al. 2019 “Achieving Conservation of Energy in Neural Network Emulators for Climate Modeling” (arxiv.org/abs/1906.06622). In that paper, two options to implement physics constraints were discussed: 1. Constraining the loss function; 2. Constraining the architecture. The approach of this paper belongs to option 2 (figure 2, Beucler et al. 2019). It presents a detailed derivation and implementation of the constrained layer with fixed weights for a photochemical model.

Line 105, it is stated that a key difference with the work of Beucler et al. 2021 is “1. Our entire output is calculated under the constraints, rather than only a portion of the output.” In my opinion, this statement is not correct. In the general architecture (Beucler et al.), parts of the output can be constrained and some parts of the output don’t need to be constrained. It’s a choice, to constrain the entire output is just a particular case.

The authors compared their constrained NN with a simple one-layer NN (naïve NN). The constrained NN has two layers (though the 2^nd layer has fixed weights). A two-layers NN could (possibly) improve the performance and it is also more comparable to their constrained NN, it would be fair to compare the constrained NN with a 2-layers naïve NN.

Some minor textual comments:

• (1), use a different symbol for the function C (C already refers to concentrations);
• (2), use a different symbol for function S (the 2^nd S)
• Line 221, 1.2 million samples should be: 0.12 million samples
• Line 232, “126 days with continuous observations”, data in each day are generated/initialized separately (see Line 220), so they are not continuous?
• Line 240, Svector à S vector