Neural network classification of eigenmodes in the magnetohydrodynamic spectroscopy code Legolas

To predict the immediate evolution of a plasma system, one needs to identify the nature of the dominant instability. In this work, a neural network is employed to address a non-binary classification problem of instabilities in astrophysical jets, whose natural oscillations and instabilities are quantified with the magnetohydrodynamic spectroscopy code Legolas. The trained models exhibit reliable performance in the identification of the two instability types supported by these jets. To improve the neural network aided classification process, techniques for training data augmentation and refinement of predictions for general eigenproblems are discussed.


Introduction
For many plasma-or fluid-related disciplines the question of stability is of central interest.In fusion research instabilities break confinement (Nuhrenberg et al, 1993), in solar physics they lead to eruptions like coronal mass ejections (Lynch et al, 2016), and in space weather they affect the propagation and properties of the solar wind (Shaaban, S. M. et al, 2019).Since many plasma configurations can host a plethora of instabilities, determining the driving forces or physical effects which give rise to them is crucial.Additionally, when several instabilities coexist, it is essential to identify which one is dominant, i.e. grows most rapidly, and thus governs the initial evolution of the plasma configuration.
To address this central question, the magnetohydrodynamic (MHD) spectroscopic code Legolas (Claes et al, 2020, and https://legolas.science)was developed, which allows for the investigation of the influence of physical parameters, such as flow and resistivity, on the dynamics of a plasma configuration, and in particular, on its magnetohydrodynamic stability.For a given equilibrium structure and a choice of non-ideal effects (i.e.resistivity, viscosity, etc.), the Legolas code computes all the waves and instabilities supported by the plasma (also referred to as modes).Using this tool, one can obtain a comprehensive overview of the instabilities and their respective growth rates within the considered parameter space.However, the results also contain a multitude of modes which are not necessarily relevant to stability, like sequences of slow, Alfvén, and fast waves.Consequently, it may become difficult to track any specific mode during the exploration of the parameter space or pinpoint which effect is causing it.Hence, we seek an algorithm for categorising modes, in order to distinguish and classify the relevant instabilities.
The process of classification, where one assigns a label from a finite, predefined set to an object, is a notoriously time-consuming task if performed manually.Hence, it is desirable in many fields to automate the classification of data.With the continuous development of new machine learning techniques, many architectures have been explored for classification (Zhang, 2000;Kiranyaz et al, 2021;Moosaei and Hladík, 2023).Here, we improve on the results of Kuczyński et al (2022), introducing a supervised neural network designed for the non-binary classification of any generalised eigenvalue problem.We apply the model to the study of an astrophysical jet (Belan et al, 2013), which are also replicated in recent experiments (Bellan, 2018), here with shear axial flow embedded in a helical magnetic field (Baty and Keppens, 2002), and demonstrate reliable performance.
First, we introduce the equations solved by the Legolas code in Sec. 2, and then describe the particular physical system used for testing the model in this work.In Sec. 3, we describe the eigenvalue classification algorithm, and present a method for enlarging training data and refining model predictions.Subsequently, Sec. 4 outlines how the described algorithm is applied to the test problem, focusing on data preparation, suggested neural network architectures, overall performance metrics, and the filtering criterion for 'uninteresting' modes.Finally, in Sec. 5 we verify the performance of the algorithm, comparing the outcome between the two introduced network architectures.

Astrophysical jets in the Legolas code
Before introducing the neural network based algorithm, we briefly describe the data generated with the Legolas code.The MHD spectroscopic code Legolas (Claes et al, 2020;De Jonghe et al, 2022;Claes and Keppens, 2023) is a finite element method (FEM) code that solves the generalised eigenproblem that arises after linearisation and 3D Fourier analysis of a set of (magneto)hydrodynamic equations.For the data used in the present classification problem, the equations are linearised around an equilibrium representing an astrophysical jet with shear axial flow embedded in a helical magnetic field, as described by Baty and Keppens (2002).The equilibrium is described by a constant density ρ 0 and velocity, magnetic field, and temperature profiles where V is the asymptotic velocity, R j the jet radius, a the radial width of the shear layer, r c the characteristic length of the radial magnetic field variation, B θ and B z magnetic field strength parameters, and T a the temperature at the jet axis.For this study, 240 Legolas runs of this configuration were carried out in the interval r ∈ [0, 2] for various values of V .At r = 0, a regularity condition was imposed, whilst a perfectly conducting boundary condition was used at r = 2. Table 1 gives an overview of all the parameter values.
In each Legolas run, the resistive, compressible MHD equations, perturbed and linearised around the equilibrium (2-4), are solved for the frequency ω and Fourier amplitudes f1 (r) after substituting a Fourier form, for each perturbed quantity (ρ 1 , v 1 , T 1 , B 1 ), with imposed wave numbers m and k.In these MHD equations, p represents the pressure, governed by the ideal gas law p = ρT , and J = ∇ × B the current.Furthermore, η is the resistivity, set to η = 10 −4 , and γ the adiabatic index.Since we employ a fully ionised, non-relativistic approximation, the adiabatic index is set to γ = 5/3.For a grid discretisation with N grid points, Legolas computes 16N complex eigenvalues ω and their 8 corresponding (complex) eigenfunctions ρ 1 , v 1 , T 1 , and B 1 on a grid with 2N − 1 grid points.Hence, for the classification algorithm, each of the resulting 16N data points is treated as a 2-vector containing the real and imaginary parts of the eigenvalue along with its corresponding complex 8 × (2N − 1) matrix containing the eigenfunctions.
For this jet configuration, the associated spectrum of complex eigenvalues contains up to two types of instabilities: one Kelvin-Helmholtz instability (KHI) and a parameter-dependent amount of current-driven instabilities (CDI) (Baty and Keppens, 2002;Hardee, 2011).The spectra for two distinct parameter choices are shown in Fig. 1(a,b) as examples.In Fig. 1(a) the KHI and the sequence of CDIs are indicated.Though they are easily identifiable in this first case by their position in the spectrum, this is harder for the case in Fig. 1(b), where some modes are not fully resolved at this resolution.In general, we identify the instabilities by their eigenfunction behaviour.The real part of the ρ-eigenfunction of the KHI, visualised in Fig. 1(c), is characterised by a maximum located at the jet boundary (r = R j ) whereas CDIs are characterised by a smooth, oscillatory behaviour inside the jet (r < R j ), as illustrated in Fig. 1(d).Modes that do not possess any of these characteristics are referred to as uninteresting.

Class preserving maps
The goal of classification algorithms is to associate a unique label l ∈ L with an input x ∈ X, where L and X are sets.Mathematically, this is a function from X to L, i.e., Class : X → L. (10) In the present work, this map is realised via a supervised neural network and a subsequent, user-informed filtering procedure.In order to gain a better understanding of the structure of this algorithm, in what follows, we introduce the concept of class preserving maps.
We define the training dataset T ⊂ X that consists of all training data points t ∈ T such that the map in Eq. ( 10) is known.Consider the set of maps U from X to itself that preserve the class label.Denoting such a map by u, we have It is apparent that Class(u(t)) = Class(t).If u is an injective map satisfying u(T ) ⊂ X \ T , any element u(t) can be used to extend the training dataset T .However, if u is not injective or u(T ) ∩ T ̸ = ∅, care should be taken not to include repeated elements in the dataset, which could introduce an imbalance in the training data.Additionally, let x ′ ∈ X \ T be an input whose class label must be inferred by the neural network.Rather than making a prediction on a single input x ′ , one can also compare it with u(x ′ ) which, in an ideal scenario, should result in the same label.The user is then able to choose a prediction dependent on their preferred filtering scheme.If the model is free from systematic errors, this results in a higher likelihood of correct classification.However, if a certain map u l only preserves the class of a subset X l ⊂ X, it cannot be used for the reinforcement of the network's prediction, since, a priori, we do not know if the class of the data point being predicted will be unaltered.Nevertheless, u l can still be used for data generation.These two types of class preserving maps are illustrated in Fig. 2.

Structure of the eigenvalue classification algorithm
In some applications, the input to the neural network is an ordered tuple.This could be, for example, a text-image pair or, as in our problem, an eigenmode-eigenfunction pair (ω, f ω ), where the subscript ω now indicates that the eigenvector f ω is associated with the eigenvalue ω.In such scenarios, it is common to implement separate branches for different constituents of the input (e.g.Boldeanu et al, 2021).In our model, we first extract convolutional features of f ω in a separate branch, which results in a reduced representation f ω .Then, ω is simply concatenated with f ω .The combined result (ω, f ω ) is then further fed into a regular neural network that in the end returns the probability of each class label l.Then, probability thresholds are optimised in order to maximise the chosen metric which judges the performance of the model.Finally, once the neural network is trained and the thresholds are chosen, a filtering scheme is incorporated based on the previously defined maps u i ∈ U .The complete scheme of the eigenvalue classification algorithm as applied to the KHIs and CDIs classification problem is shown in Fig. 3.

Application to Legolas jet data
Here, we apply the algorithm described in Sec. 3 on the KHI-CDI classification problem introduced in Sec. 2. First, we return to the data structure, and discuss how the data is expanded with the use of class preserving maps.Then, we describe the network architecture presenting the network's layers in more detail.Subsequently, we discuss the chosen performance metrics and how we optimised them by introducing probability thresholds and a filtering procedure.
2. An eigenfunction input as a complex matrix f ω of dimensions 301 × 8. Its real and imaginary parts are stored in 2 channels.
We decided to use 80% (192) of these runs for training, 10% (24) for validation, and 10% for testing.The division of the files across the three categories was randomised.
For the exact distribution, see Table 4 at the end.The resulting data contained 99.76% uninteresting modes (class 0), i.e. modes that are neither KHI (class 1) or CDI (class 2).Therefore, in order to balance and extend the dataset, we utilised the technique of class preserving maps, described in Sec.3.1.We defined the following maps: 1. Multiplying the eigenfunctions by a complex phase factor: This map u ∈ U is always class preserving because eigenfunctions are only determined up to a complex factor.As a consequence of this property, we can also use these maps to decrease the uncertainty of the neural network's prediction in the filtering step, as discussed in Sec.3.1.2. Data superposition: if θ = z = t = 0 in Eq. ( 9), the corresponding eigenvalueeigenfuction tuples satisfy the following superposition principle, with ω = ω ′ and φ, φ ′ ∈ (0, 2π).Nevertheless, when employed for ω ̸ = ω ′ , this map typically preserves the inherent characteristics of the considered modes, i.e., the peak at the jet boundary (r = R j ) for KHI modes, and the oscillatory behaviour inside the jet (r < R j ) for CDI modes.Additionally, the purpose of the sum 1 2 (ω + ω ′ ) is to artificially assign information about the growth rate to the created mode.Therefore, we treated Eq. ( 13) as an approximate class preserving map for general modes of class l.
The resulting distribution of initial, and final training, validation, and testing data is summarised in Table 2.
Table 2 Distribution of initial and generated data samples across different classes for neural network classification training, validation, and testing phases.For each class, the values are given as a percentage (%) of the total number of modes in the last column.Entries associated to u and u l are then the percentage generated using Eqs.( 12) and ( 13

Network architecture
We propose two different neural network architectures, one for high performance computers and one for single thread computations.The former is a variant of the ResNet (He et al, 2016) and the latter is a plain network.The models were developed in Keras (Chollet et al, 2015), an open-source neural network library written in Python.First, we discuss the architecture of the ResNet.As described in Sec.3.2, the network consists of two stages.Initially, only the input eigenvector f ω is passed through a convolutional feature extractor.Then, the result f ω is concatenated with the eigenvalue ω and further fed into the second stage.Fig. 4 shows the main building blocks of the network used in both stages, where some layers are marked with a symbol for reference here.In the first stage, the weights layers (*) are convolution layers whilst in the second stage they are dense layers.The kernels of the convolutional layers within a building block are of the same size.Both stages consist of three such main building blocks.The kernel sizes are (65, 5), (33, 3), (17, 2) and the number of convolutional filters is 128, 64 and 32 accordingly.The intermediate dense layers in the second stage have sizes 34 and 17.The dropout ( †) value is 0.1 for convolution layers and 0.5 for dense layers.Average pooling ( ‡) is only used for convolution layers.We chose the Adam optimizer with a learning rate of 0.001, β 1 = 0.9, and β 2 = 0.999.The loss function is the categorical cross-entropy.
Regarding the plain network, the main difference is that it follows the skip connection only ( § top), and the main branch ( § bottom) is removed.Since the training was much more computationally inexpensive than in the case of the ResNet, we were able to fine-tune the network's hyperparameters more.In fact, the kernel sizes, number of filters, and intermediate dense layer sizes listed in the previous paragraph were chosen as such, since these were the ones that were close to optimal for the plain network.

Probability thresholds and performance metric
In this application, the primary objective of the eigenfunction classification algorithm is to identify the relevant modes from a large dataset, which contains predominantly uninteresting modes.The final model should only misclassify a few relevant modes as uninteresting whilst correctly filtering out most of the truly uninteresting modes.Hence, for evaluating the model's performance, we chose two metrics: precision and (17 This way, we emphasise the priority of identifying relevant modes over distinguishing between classes 1 and 2. Finally, as an informative metric, we introduce the balanced accuracy, balanced accuracy = 1 2 After training the network, the predicted probabilities for each class are denoted as (p 0 , p 1 , p 2 ).In the inference step, instead of simply selecting the class with the highest probability, we introduce thresholds a and b.We first determine if p 1 ≥ a.If it is, class 1 is predicted.Otherwise, we evaluate if p 2 ≥ b.If this condition is met, class 2 is predicted.Otherwise, class 0 is the default prediction.To improve the classification algorithm, the thresholds a and b are optimised using validation data.This is done by imposing a desired recall value, i.e. a tolerance on the number of discarded relevant modes.We then seek the highest possible precision as a function of a and b.Since the distributions of validation and testing data are different (Table 2), during testing we can expect the recall and precision values to deviate from their validation values, that are the results of this optimisation procedure.

Filtering
Once the thresholds for a and b are established, the eigenfunctions of the testing data can again be subjected to the class preserving maps of the form (12).The resulting data couples (ω, e iφ f ω ) are evaluated by the network and classified according to the optimised thresholds a and b.Repeating this for m different phases φ k (k = 1, . . ., m) results in a total of m + 1 predictions (including the unmodified data).Subsequently, the final label is the label that was predicted the most often, with ties broken in favour of relevant over uninteresting, and if decidedly relevant, class 1 taking precedence over class 2. This technique could be further improved by considering the variance of the predictions in line with the work of Gal and Ghahramani (2015).Furthermore, both techniques could be used simultaneously.Nevertheless, we achieved satisfactory results with calculating only the most likely prediction from the ensemble of predictions generated from class preserving maps which we report in the next section.

Results
First, the plain network was trained on 900 batches of 512 modes each.Then, using validation data, probability thresholds a and b were optimised such that the corresponding validation recall was at least 90%.Next, five predictions were generated, and a final label was extracted, as described in Sec.4.4.The resulting confusion matrix is shown in Fig. 5(a).The network thus achieves a recall of 94.3%, a precision of 36.3%, and a balanced accuracy of 97.0% on the testing data.Therefore, from the modes classified as 1 or 2, 36.3% were truly relevant modes, whilst 5.7% of all relevant modes were lost.
Similarly, the ResNet was trained on 2230 batches of 512 modes each, but the a and b thresholds were optimised for a minimal validation recall of 95%.Following the same classification process as the plain network resulted in the confusion matrix in Fig. 6(a).Thus, the ResNet reached a recall value of 97.6%, precision of 38.0%, and a balanced accuracy of 98.7% on the testing data, outperforming the plain network in all metrics.
By imposing different minimal validation recall values, we can control how many relevant modes are lost.Of course, adapting the validation recall also changes the testing recall, precision, balanced accuracy, and thresholds (a, b).For varying validation recall values, the testing recall, precision, and thresholds are visualised in Fig. 5(b) for the plain network and Fig. 6(b) for the ResNet.Unsurprisingly, higher validation recall values lead to higher testing recall, and lower testing precision and thresholds.As the graphs show, the final recall is consistently higher than the imposed validation recall, with a validation recall of 95% sufficing to achieve a perfect recall for the plain network (at the cost of lower precision).It is remarkable however that the a threshold remains constant at a high value of 71% for varying recall values.This implies that the network assigns high probabilities to class 1 if the mode is truly a class 1 mode.It also means that if the minimal validation recall is increased to 95%, there are no additional modes mislabelled as class 1, and the decrease in precision is solely due to the misclassification of class 0 modes as class 2 modes.
For the ResNet this behaviour is even more pronounced.Since its a value is close to 1, with the b value only slightly lower, the network has to assign extremely high probabilities to either class 1 or 2 to consider a mode relevant.Additionally, imposing a validation recall larger than 95% results in vanishing a and b values such that everything is classified as class 1.This indicates that the three lost relevant modes in the lower left of Fig. 6(a) are never classified as relevant unless all other modes are too.
Returning to the confusion matrices, two more observations are noteworthy.Firstly, the lower right 2 × 2 submatrix is diagonal for both networks.Hence, they clearly distinguish between class 1 and class 2 modes.This is in line with initial expectations, based on the eigenfunction shapes, like those shown in Fig. 1(c,d), which show strong behaviour at the jet boundary for KHI modes in contrast with the CDI behaviour in the jet's interior.Secondly, the relative number of class 0 modes that were misclassified as class 1 is much smaller than the number misclassified as class 2 for the plain network.In addition, it is worth noting that all class 1 modes were correctly identified by both networks.For the networks in Figs.5(a) and 6(a), all metrics have been computed per class and are displayed in Table 3. From this table it is clear that the ResNet has better recall values, and achieves a higher precision in class 2 at the cost of a lower precision in class 1.In both cases, the network is better at identifying modes of class 1 than of class 2.

Conclusion
Due to the large amount of natural oscillations of a single plasma configuration (one run of the MHD spectroscopic code Legolas), visual inspection of the eigenmodes to identify characteristics can be a monotonous and time-consuming task.In this work we have applied two convolutional neural networks to a non-binary classification problem of ideal MHD eigenmodes in astrophysical jets, analysed with Legolas.For a recall of 94.3% the plain neural network left 0.55% of all modes for manual inspection.The ResNet offered a significant improvement with a recall of 97.6% leaving 0.43% for inspection.Furthermore, neither network confused class 1 with class 2 modes in the test data.Since even the plain network provided good results already, we conclude that neural networks offer a great opportunity for automated mode detection in Legolas data, and likely, for classification of eigenproblem data in general.
In the context of large parameter studies with the Legolas code, these results are particularly promising.With an automated way of reliably identifying instabilities, various parameters can be varied simultaneously, and the resulting parameter space can be partitioned into sections of similar behaviour efficiently.This approach could have many applications, e.g. in the analysis of jet stability like the problem presented here, or the problem of current sheet stability in various astrophysical settings, where tearing and Kelvin-Helmholtz instabilities compete (Hofmann, 1975;Li andMa, 2010, 2012).
A significant drawback of the supervised approach however is of course the need for a large set of pre-classified data for training purposes.To sidestep this issue, future investigations could focus on unsupervised clustering algorithms to search for structures in Legolas data, like the translational-azimuthal distinction in Taylor-Couette flows (Dahlburg et al, 1983) or the surface-body wave dichotomy in flux tubes (Edwin and Roberts, 1983).For large parameter studies like those described in the previous paragraph, this method is especially appropriate, considering it requires less prior knowledge of the types of modes one may encounter throughout the parameter space.
Finally, it remains an open question whether a generally-applicable neural network for Legolas data is possible.In particular, is it feasible to develop a neural network that can predict which physical effect, like shear flow or magnetic shear, is responsible for each instability in a spectrum?In this regard, another hurdle to overcome is that

Fig. 2
Fig. 2 Two types of class preserving maps.The map u ∈ U , indicated by the blue solid line, preserves the class of all eigenfunctions f ∈ X.The map u 1 ∈ U 1 , indicated by the red dashed line, preserves only the internal maps of the subspace denoted as X 1 in the figure.

Fig. 3
Fig. 3 A schematic of the employed classification algorithm.

Fig. 4
Fig. 4 Diagram of the main building blocks of the neural networks.
negatives are modes that the model labelled as relevant (class 1 or 2) and uninteresting (class 0), respectively.Denoting the elements of the confusion matrix as c ij , where the rows correspond to the true label, and the columns to the predicted label, the precision and recall are given as precision = d 11 d 01 + d 11 , recall = d 11 d 10 + d 11 , (16) with d ij = d 00 = c 00 , d 01 = c 01 + c 02 , d 10 = c 10 + c 20 , d 11 = c 11 + c 12 + c 21 + c 22 .

Fig. 5
Fig. 5 Plain network.(a) Confusion matrix for a minimal validation recall of 90%.(b) Recall, precision, and thresholds as functions of the imposed validation recall.

Fig. 6
Fig. 6 ResNet.(a) Confusion matrix for a minimal validation recall of 95%.(b) Recall, precision, and thresholds as functions of imposed validation recall.

Table 1
Parameters of the data used in this study.The upper table shows the different values of V in the dataset.For each value of V , k was varied from 0.5 to 7 in increments of 1/6.The parameters in the lower table were identical in all cases.

Table 3
Precision and recall (%) per class for both the plain network and the ResNet.

Table 4
Files in the dataset used for validation and testing.The remaining files were used for training.