Exploiting a variational auto-encoder to represent the evolution of sudden stratospheric warmings

Sudden stratospheric warmings (SSWs) are the most dramatic events in the wintertime stratosphere. Such extreme events are characterized by substantial disruption to the stratospheric polar vortex, which can be categorized into displacement and splitting types depending on the morphology of the disrupted vortex. Moreover, SSWs are usually followed by anomalous tropospheric circulation regimes that are important for subseasonal-to-seasonal prediction. Thus, monitoring the genesis and evolution of SSWs is crucial and deserves further advancement. Despite several analysis methods that have been used to study the evolution of SSWs, the ability of deep learning methods has not yet been explored, mainly due to the relative scarcity of observed events. To overcome the limited observational sample size, we use data from historical simulations of the Whole Atmosphere Community Climate Model version 6 to identify thousands of simulated SSWs, and use their spatial patterns to train the deep learning model. We utilize a convolutional neural network combined with a variational auto-encoder (VAE)—a generative deep learning model—to construct a phase diagram that characterizes the SSW evolution. This approach not only allows us to create a latent space that encapsulates the essential features of the vortex structure during SSWs, but also offers new insights into its spatiotemporal evolution mapping onto the phase diagram. The constructed phase diagram depicts a continuous transition of the vortex pattern during SSWs. Notably, it provides a new perspective for discussing the evolutionary paths of SSWs: the VAE gives a better-reconstructed vortex morphology and more clearly organized vortex regimes for both displacement-type and split-type events than those obtained from principal component analysis. Our results provide an innovative phase diagram to portray the evolution of SSWs, in which particularly the splitting SSWs are better characterized. Our findings support the future use of deep learning techniques to study the underlying dynamics of extreme stratospheric vortex phenomena, and to establish a benchmark to evaluate model performance in simulating SSWs.


Introduction
Sudden stratospheric warmings (SSWs) are the primary dynamical extreme events observed in the wintertime stratosphere, with all but one observed SSW having occurred in the Northern Hemisphere (Baldwin et al 2021).Specifically, a major SSW is typically associated with stratospheric polar-cap temperatures rising 30-40 K in a few days and a distorted polar vortex accompanied with, in the zonal-mean sense, a reversal of polar stratospheric westerlies (Andrews 1987, Limpasuvan et al 2004, Charlton and Polvani 2007).Theoretical and modeling frameworks to understand the dynamics of SSWs are mainly based on linear wave theory, involving the interaction between the zonal flow of the stratospheric polar vortex and the upward propagation of planetary waves (Charney and Drazin 1961, Matsuno 1971, Schoeberl 1978, been explored.This motivates our study to use deep learning models directly on SSWs, in order to provide a new perspective on their spatiotemporal evolution. To tackle the issue of the small number of observed SSWs (i.e.identified in reanalyzes), large numbers of SSWs simulated by state-of-the-art global climate models with large ensembles could be exploited.This idea is conceived from recent deep learning studies investigating global warming and climate variability patterns, in which large amounts of global temperature longitude-latitude maps were sampled from multiple CMIP5/6 repositories and then used to train a non-linear deep learning model (e.g.Barnes et al 2019, Ham et al 2019, Gordon et al 2021, Labe and Barnes 2021, 2022a, 2022b, Gordon and Barnes 2022).Global climate models have been widely used to strengthen our process-based understanding of stratosphere-troposphere coupling during SSWs (Ayarzagüena et al 2019, Baldwin et al 2021).Simulated SSWs, in particular from the so-called 'high-top' models, possess similar characteristics to observed SSWs in terms of their occurrence frequency, spatial extent, and temporal evolution (Charlton-Perez et al 2013, Ayarzagüena et al 2019).Thus, if one can sample large numbers of SSWs from large-ensemble climate model simulations as the training dataset, these SSWs would allow us to train a deep learning model in a more feasible way.
As each SSW manifests a predominant spatial pattern in the stratosphere (either displacement-type or split-type structure), deep learning models, especially those that are salient in extracting features from two-dimensional images, could offer a new opportunity to examine SSW morphology.Convolutional neural networks (CNNs) are a deep learning technique developed to tackle such image recognition tasks and have been widely used in multiple earth science problems to extract key spatial structures (Ham et al 2019, 2023, Connolly et al 2023).In addition, variational auto-encoders (VAEs; Kingma and Welling 2013), a typical type of generative deep learning model, have been extensively utilized in retrieving non-linear relationships in geoscience data (Vuyyuru et al 2021, Lopez-Alvis et al 2022, Tsekhmistrenko et al 2022).VAE models consist of two components: an encoder and a decoder.The encoder maps the input data into a latent space, which represents a compressed form of the data, whereas the decoder samples the compacted representation from the encoder under given constraints to generate corresponding output data.Combining CNN and VAE results in a deep learning model that can not only be employed for image or pattern recognition tasks but also can efficiently generate continuous transitions of pattern structures from one state to another.Therefore, a joint CNN-VAE model could offer an innovative means to construct a phase diagram, allowing a projection of the SSW spatiotemporal evolution onto it, yielding insight into the physical and temporal evolution of SSWs.Thus far, Krinitskiy et al (2019) is, to the best of our knowledge, the only deep learning study that trained a CNN-VAE model and analyzed the latent space with agglomerative clustering to investigate stratospheric circulation variability.
The goal of the present study is to explore the capability of CNN-VAE, inspired by Krinitskiy et al (2019), in representing the spatiotemporal evolution of SSWs.For simplicity, we refer to CNN-VAE as VAE throughout the manuscript.We take advantage of a suite of large-ensemble historical simulations carried out with the Whole Atmosphere Community Climate Model version 6 (WACCM6; Gettelman et al 2019, Liang et al 2020, 2022) to sample a large number of SSWs as the training dataset.The novelty of this study is to use SSWs from WACCM6 to train a VAE model and to create a representation to map the spatiotemporal evolution of SSW events in the phase diagram derived from the latent space.
The remainder of the paper is laid out as follows.In section 2, we first introduce the reanalysis data and WACCM6 results used in this study and elaborate on the structure and training process of VAE.In section 3, we present and discuss the results of VAE and compare them with those obtained from PCA.We also demonstrate that VAE trained with modeling results can be applied to SSWs from reanalysis.In section 4, we discuss the physical interpretation, improvements, and applications of the VAE phase diagram, followed by a conclusion in section 5.

Reanalysis data
We use daily-mean zonal wind and geopotential height fields at 10 hPa (hereafter U10 and Z10, respectively) from the National Aeronautics and Space Administration's Modern-Era Retrospective Analysis for Research and Applications, version 2 (MERRA-2; Gelaro et al (2017)) with a 0.5 • latitude × 0.625 • longitude grid.We use the MERRA-2 SSWs as the observational samples to test if the capability of VAE trained with simulated SSWs can be applied to the observed SSWs.

Global climate model simulations with large ensembles
In this study, we employ two sets of historical simulations over 1979-2014, conducted using WACCM6 (Gettelman et al 2019).WACCM6 is the high-top version of the atmospheric model within the Community Earth System Model version 2 (CESM2; Danabasoglu et al 2020), with 70 levels extending from the surface to the model upper boundary at 6 × 10 −6 hPa (∼140 km).The horizontal resolution of WACCM6 is 0.95 • and 1.25 • in longitude and latitude respectively.While WACCM6 shares identical parameterizations for most physical processes with CAM6 (Danabasoglu et al 2020), there exists a difference in the parameterization scheme for gravity waves (Gettelman et al 2019): although both models employ identical orographical gravity wave drag parameterizations, distinct parameter values were incorporated for the frontal and convective parameterizations in WACCM6, leading to improved representations of frontogenesis.Compared to CAM6, WACCM6 has demonstrated superior performance in simulating stratospheric processes including SSWs, tropospheric circulation variability (e.g.storm tracks, stationary waves, and midlatitude blocking), stratosphere-troposphere coupling, and a self-generated Quasi-Biennial Oscillation (e.g.Ayarzagüena et al 2019, Gettelman et al 2019, Simpson et al 2020).
The first set of simulations involves WACCM6 being forced by the daily time-varying global sea-surface temperature (SST), sea-ice concentration (SIC), and CMIP6 historical radiative forcings (Eyring et al 2016, Haarsma et al 2016).This configuration encompasses the comprehensive effects of all observed forcings on the atmosphere and is usually referred to as the Atmospheric Model Intercomparison Project-type experiments.The second set of simulations replicates the first set with the exception that the daily time-varying SIC field for the Northern Hemisphere is replaced by the daily 1979-2014 climatological values.This alteration ensures that the atmospheric circulation in the second set of simulations remains unaffected by variations in Arctic sea ice.The prescribed boundary conditions for SST and SIC between 1979 and 2014 are obtained from the Met Office Hadley Centre Sea Ice and SST version 2.2.0.0 dataset (Kennedy et al 2017), which are also utilized in the CMIP6 HighResMIP protocol (Haarsma et al 2016).To ensure consistent evolution of the SST and SIC fields in both simulation sets and to eliminate any non-physical values in these fields, we carried out a SIC-SST adjustment in accordance with Hurrell et al (2008) (see section 2.1 of Liang et al (2020) for further details).It is noted that several dynamical and statistical aspects involving stratospheric circulation variability and stratosphere-troposphere coupling are very similar in the two experiments, indicating minimal SIC impacts on the stratosphere.
In each case, we perturb the initial atmospheric temperature fields with small values (the so-called pertlim or micro-perturbation method; Kay et al 2015, Deser et al 2020), to construct a total of 30 members for each set of simulations, giving a total of 2160 simulation years (2 sets of simulations times 36 year integration times 30 ensemble members).Such large-ensemble simulations yield a total of 1539 SSWs (see section 2.3 for the description), more than 35 times the number of observed events.Detailed comparisons and discussion between the SSWs simulated in WACCM6 and those detected in MERRA-2 and other reanalysis datasets are included in ongoing work.

Identification of SSWs
To detect SSWs in both reanalysis and WACCM6 simulations, we follow the criteria proposed by Charlton and Polvani (2007).This widely-used algorithm detects SSWs based on the reversal of the zonal-mean U10 at 60 • N from westerly (i.e.positive) to easterly (i.e.negative) in boreal extended winter, with additional criteria to separate SSWs and distinguish between SSWs and the final warming in spring.The onset or central date, which we denote as 'day 0' , is then determined as the day on which the zonal winds reverse.Then, 'day −10' denotes the date ten days prior to the onset date, and 'day 10' the date ten days after the onset date.Square brackets are used to denote the period of SSW evolution.For example, [−10, 4] means the period of day −10 to day 4. We identify 23 SSW events during the 1979-2014 period in MERRA-2 and 759 and 780 events in the first and second sets of WACCM6 simulations, respectively.We also compared the results from the latest European Centre for Medium-Range Weather Forecasts reanalysis (ERA5; Hersbach et al (2020)) and find the same results in terms of SSW occurrence frequency.
To distinguish between displacement and split-type SSWs, we modify the method proposed by Seviour et al (2013), which uses only Z10 field (hereafter, Z10 method).We apply this method to the Z10 spatial patterns of the identified SSWs, which results in the two-dimensional vortex moment to specify the vortex structure during days [−10, 20].Two key quantities are determined: the centroid latitude (ϕ o ) and the aspect ratio (r).A displaced vortex is defined when the centroid latitude remains equatorward of 66 • N for at least 7 d, while a split vortex requires the aspect ratio to remain higher than 2.4 for at least 7 d.If these two criteria are not satisfied, we categorize the event as displacement-type unless none of the centroid latitudes are south of 66 • N and the aspect ratio is larger than 2.4 for at least 1 d within 14 d after the onset date.

Training data preprocessing
The Z10 field output from WACCM6 simulations are in a nominal 0.95 • × 1.25 • horizontal resolution with uniform grid spacing.In order to maintain the grid area size at high latitudes, we employ a bilinear interpolation method (Kim et al 2019).This spatial interpolation also allows the VAE model to extract features in a more feasible way (i.e.squared images) as the stratospheric polar vortex can be distorted in the longitude-latitude grid.Figure 1(a) exhibits the original longitude-latitude grid in an orthographic projection with 90 • N at the center.The region of color shading indicates the domain chosen in this study.It is apparent that the grid area becomes smaller toward the north.After using the bilinear interpolation method, we construct 36 × 36 squared images with a similar grid size over the selected domain (figure 1(b)).Before training, these squared images of Z10 fields are then standardized to the Z-score to accelerate the convergence rate (Ioffe and Szegedy 2015): where µ X and σ X denote the mean and the standard deviation of the whole set of Z10 fields, x ij is a particular Z10 image pixel at grid location i and j.We prepare a total of 32 319 Z10 images from the Z10 anomalies over 10 d prior to and after the onset date of 1539 SSWs, in which 1177 events are displacement-type and 362 events are split-type.We adopt a 80%-9%-11% ratio of data for separating training, validation, and test datasets, resulting in 25 893 Z10 images for training, 3024 images for validation, and 3402 images for testing.
To incorporate spatiotemporal evolution information, we classify the daily aspect ratio (r) > 2.4 as the 'splitting pattern' , centroid latitude (ϕ o ) < 66 • N as the 'displacement pattern' , and all others as 'neither' .To avoid sampling bias toward one type of event, we randomly choose 1299 images from the each pattern, yielding a balanced training dataset of 3897 images.The number 1299 is the count of images of splitting pattern, which has the smallest size of images among three patterns.

Variational auto-encoder (VAE)
The VAE (Kingma and Welling 2013) is a generative model composed of an encoder neural network and a decoder neural network (figure 2).The encoder propagates the input data to a compressed representation, which is called the latent space.The decoder samples the compacted representation by the given conditions (usually a Gaussian distribution) and generates the output data.The VAE is trained by minimizing the difference between the output and input data constrained by a Gaussian distribution in the latent space, optimizing the latent space to represent the continuous variations of the patterns from the input data.
Looking into the latent space yields an understanding of how the VAE learns the continuous variations in the patterns.We can choose how many representations in the latent space are used in a VAE and thus obtain the 'modes' or 'dimensions' of the latent space.Due to the nonlinearity and complex VAE network, these modes may provide an alternative way to represent the structure of the phenomenon that may not be resolved by linear data analysis methods, such as PCA.However, there is no consensus or theory on how many 'modes' should be used.In this study, we are motivated by the typical diagnosis of the Z10 pattern with two variables (e.g.r, ϕ o , cf section 2.3 or Seviour et al (2013)) to describe the spatial distribution and evolution of the stratospheric polar vortex.Thus, we focus on the two 'modes' of the trained VAE.Moreover, two modes or dimensions can form a two-dimensional phase diagram, making it easy to examine the evolution of the SSWs.
Figure 2 illustrates the VAE structure and its network size used in this study.The second row depicts the encoder using a 3 × 3 kernel.The first few convolution layers extract features to 32, 64, and 64 channels, preserving the 36 × 36 pixels.The following layers transform to 18 × 18 pixels with 64 channels and 9 × 9 pixels with 128 and 256 channels, respectively.A dense layer with 64 neurons is followed to concatenate information to two pairs of the means and log-variances: (µ 0 , v 0 ) and (µ 1 , v 1 ), which are passed to the latent space with two representations or sampled results (z 0 , z 1 ): where N(0,1) is a Gaussian distribution with zero mean and unit standard deviation.The third row shows the decoder, which takes the sampled results z i as inputs.The decoder reverses the data-convolving process of the encoder to generate the output data with a size of 36 × 36 pixels.The top row summarizes the decoder and encoder connecting with a latent space (light blue color shading).It is noted that the vertical solid lines after convolution layers denote the max-pooling/upsampling layers, which will reduce/increase the pixels.The loss function is determined by the mean squared error (MSE) between the input and the generated output, plus the Kullback-Leibler divergence between the latent space and the Gaussian distribution.The backpropagation algorithm is performed to calculate the derivatives along with the Adam optimizer to minimize the loss.
To visualize the latent space, we utilize the decoder of the trained VAE.We select varying values of z 0 and z 1 , from −4 to 4, in the latent space as inputs to the decoder.The outputs of the decoder give the corresponding Z10 fields with two modes or dimensions and, consequently, form a phase diagram associated with the latent space (figure 3).The reconstructed Z10 fields are obtained by inputting the original Z10 fields into the trained VAE (figures 4 and 5) and then converted to original amplitude following equation (1).For simplicity, we set the variances to be zero, although they can be non-zero.To better and quantitatively assess the performance of VAE reconstruction, we calculate the pattern correlation (R) and the root mean squared error (RMSE) between the Z10 fields of WACCM6 and reconstructed ones.

Principal component analysis (PCA)
To provide a baseline to assess the VAE results, we perform PCA over the Z10 training data.We apply a singular value decomposition to obtain the PCA modes and retrieve the spatial pattern by projecting the loading vectors onto the interpolated grid, using testing data.The first two PCA modes are utilized to construct a phase diagram with the corresponding principal components (PCs).The first PCA mode, which explains 30.6% of the total variance, is characterized by a dipole in Z10 anomalies between northwestern North America (low) and the Barents-Kara Seas (high) (figure S1(a)).The second mode, explaining 21.4% of the total variance, features lower Z10 values over Greenland together with higher values over eastern Siberia (figure S1(b)).The two modes combined explain more than 50% of the total variance and they are statistically distinguishable based on North's rule (North et al 1982) (figure S1(c)).We also tested including the third PCA mode, as it contains about 17% of the total variance, but obtained similar results (not shown).Although incorporating more PCA modes can lead to better performance (as more variance is explained), for better visualization and easier comparison to the two VAE modes, we stick to the first two PCA modes in this study.

Results
One key outcome of the VAE is the phase diagram derived from varying values of z 0 and z 1 in the latent space, which could shed insight into the morphology and spatiotemporal evolution of SSWs. Figure 3 illustrates the VAE phase diagram of SSW Z10 fields.We immediately observe that the phase diagram contains various vortex structures (Z10 < 30 000 m, depicted by blue colors).In the upper-left quadrant of the phase diagram (i.e. the regime of large negative z 0 and positive z 1 ), the vortex is rather circular and centered at or near the North Pole, with indications of a low-amplitude wavenumber-1 disturbance (stratospheric Aleutian anticyclone).Towards the lower-left corner of the diagram (i.e. the regime of negative z 0 and negative z 1 ), the vortex becomes slightly elongated and somewhat displaced toward Eurasia, but also deepens.In the upper-right corner (i.e. the regime of positive z 0 and positive z 1 ) and lower-right quadrant with smaller negative z 1 , the spatial pattern tends to manifest two smaller vortices rather than one vortex, illustrating the splitting structure of the polar vortex accompanied by an amplified wavenumber-2 disturbance.Finally, in the lower-right quadrant (the regime of positive z 0 and negative z 1 ), the vortex turns elongated without clear displacement or splitting.In general, wavenumber-1/displacement-type behavior is dominant for negative z 0 , while wavenumber-2/split-type behavior is dominant for positive z 0 .Equivalently, vortex weakening, displacement, and splitting behavior are generally characterized by positive z 1 , while less disturbed vortex states are generally characterized by negative z 1 (with the exception of the aforementioned lower-right quadrant with smaller negative z 1 ).
More intriguingly, the phase diagram exhibits a gradual, continuous transition of the vortex structure from one shape to another with varying z 0 and z 1 .For example, when z 0 is negative, the polar vortex tends to rotate clockwise, with its shape being elongated and its center displaced southward along the z 1 -axis from positive to negative z 1 values.For positive z 0 , one vortex becomes elongated when z 1 turns positive from negative and then splits into two smaller vortices as z 1 grows.Putting these together, the VAE phase diagram is capable of not only representing the important vortex features of SSWs, including displaced, stretched, and splitting shapes, but also vividly portraying their smooth transition from one form to another.To inform more details of the VAE phase diagram and the vortex transition, we show a phase diagram with finer z 0 and z 1 values in figure S2.
The PCA phase diagram also contains varied vortex structures, but some vortex shapes are not well-depicted, and their amplitude is overall stronger (figure S3) compared with the VAE counterparts.Notably, the splitting vortex structure seen in the top-right quadrant of figure 3 is not present in the PCA phase diagram, suggesting that the two PCA modes cannot appropriately portray split-type SSWs.As for the vortex transition, when PC2 is positive and large in the upper half of the diagram, the vortex rotates counterclockwise along with PC1, varying from negative to positive values.In the lower half of the diagram, when PC2 is negative, the vortex is weak or replaced by an anticyclone for negative PC1, while it is displaced towards Siberia by a strong wavenumber-1 disturbance for positive PC1.In short, the PCA phase diagram seems to be missing some important spatial features of splitting SSW vortex structure; thus, the evolution of one large vortex splitting into two smaller vortices is not expected to be well demonstrated in the PCA phase diagram using only two modes.We also show the PCA phase diagram with finer PC1 and PC2 values in figure S3.
To further shed insight on the capability of the VAE phase diagram, we first select a displacement-type SSW from the testing dataset of WACCM6 simulations and project its temporal evolution onto the diagram during the period of days [−10, 10], based on the corresponding input Z10 fields (figure 4(a)), which allows us to track its progress through the 2D phase-space.The track lies in the lower-left quadrant of the diagram, consistent with the displacement-dominant nature of the event (figure 3).The track first moves toward the origin during the precursor and onset phases of the SSW, then abruptly moves backward after nearly touching the vertical axis.This reflection of the track is consistent with the absence of any split-vortex structure (i.e.positive z 0 and negative z 1 ) during this displacement SSW, and also indicates the persistence of a fairly deep vortex core after the SSW (rather than the near-complete destruction of the vortex seen for positive z 1 and neutral z 0 ).The results indicate that the VAE has learned and generated the regime boundary, effectively separating the vortex structure of displacement-type SSWs from split-type events.We also perform the same track projection with the PCA phase diagram.The track lies within the upper-right quadrant (positive PC1 and positive PC2; figure 4(b)) and moves in a near-vertical direction from larger PC2 to smaller PC2 values.This track captures the wavenumber-1 amplification and the persistence of a strong vortex core.The PCA track and its location in the phase diagram are somewhat similar to the VAE ones, indicating their similar representation of the vortex structure and evolution.
We now examine the spatial patterns and evolution of this displacement-type SSW.The Z10 fields (from WACCM6) show that, prior to the onset date, this event features the vortex located over the North Atlantic and northern Europe.As time proceeds, the Aleutian high amplifies, and the vortex moves westward slightly and becomes elongated, with its center moving southward from 81 • N to 67 • N (figure 4(c)).After the onset date, the vortex moves eastward towards Eurasia as the Aleutian high weakens, with its center gradually turning northward.The reconstructed Z10 fields of the VAE and PCA are also displayed in figures 4(d) and (e), respectively, and they both can capture these evolving spatial features.However, differences emerge if we look into the details carefully.Before day 0, both VAE and PCA reconstructions perform similarly well: their Rs are larger than 0.92 before day −3, and RMSEs are relatively small.After the onset date, the VAE reconstructed Z10 fields perform better than the PCA reconstruction, supported by higher Rs and smaller RMSEs.The PCA reconstructed vortices are overall too circular and do not represent the stretched structure as well as the VAE reconstructed ones.
The same analysis is also conducted for a split-type SSW from the WACCM6 testing dataset (figure 5).This event features a nearly circular and strong polar vortex centered at 82 • N 10 d before the onset date, which gradually becomes elongated and begins to separate into two smaller vortices around day −6 (figure 5(c)).During the period of days [−3, 3], the vortex splitting structure is evident, with a clear wavenumber-2 configuration.During this time, the North American child vortex quickly weakens and disappears after day 3, with the Eurasian vortex becoming dominant, which is a typical evolution for a split-type SSW (Matthewman et al 2009).Afterwards, during the recovery phase, the larger Eurasian vortex begins to dominate as stratospheric wave activity diminishes.Given the more dynamic evolution of this split-type SSW compared to the selected displacement event, both VAE and PCA tracks travel all four quadrants counterclockwise (figures 5(a) and (b)) rather than within one specific quadrant as in the displacement event.A clear contrast between the VAE and PCA emerges in their reconstructed Z10 fields (figures 5(d) and (e)).While the VAE is capable of reconstructing the varied spatial patterns of the vortex-from a rather circular structure to a stretched one and to splitting patterns-the PCA fails to retrieve these important vortex features.In particular, during days [−5, 4] when the two smaller vortices appear, the PCA cannot produce the two vortices but instead only one vortex.Consequently, the Rs of PCA are overall lower than those of VAE, and the RMSEs are overall higher.
To test if the VAE, trained with SSWs from WACCM6 simulations, can be applied to observed SSWs, we select a displacement-type SSW (2 January 1987) from MERRA-2.The VAE successfully projects the evolution track, which begins in the lower-left quadrant and migrates along the z 1 -axis vertically to the upper-left quadrant (figure S5(a)).The PCA, on the other hand, gives a migration track along the vertical PC2-axis, with it lying mostly in the right quadrants (figure S5(b)).The reconstructed Z10 fields by VAE have similar Rs and RMSEs before day −3 as those by PCA, but the VAE outperforms the PCA in particular during the post-warming days [0,10].This better performance of VAE after the onset date is also shown in the WACCM6 displacement-type SSW event (figure 4(d)), suggesting that VAE is more capable of dealing with the disrupted vortex structure after the onset date.We also select a split-type SSW (6 January 2013) and find that the VAE almost always outperforms PCA during days [−10, 10] (figure S6), as for the splitting event of WACCM6 (figure 5(d)).The VAE clearly portrays the splitting structure during days [0, 5] with higher Rs and smaller RMSEs than the PCA.The above two successful Z10 reconstructions by VAE confirm that the VAE can be used to exploit the spatiotemporal evolution of observational SSWs.
We next project the Z10 fields of all displacement-type and split-type SSWs, during days [−10, 10], from the testing dataset onto the VAE phase diagram (figures 6(a) and (b)) to illustrate the holistic picture of the SSW vortex structure and to investigate if different types of SSWs favor certain regions of the 2D-phase space.To account for different numbers of events in the splitting and displacement SSW groups, we normalize by the total number of Z10 fields and show their occurrence probability as a percentage.The majority of the spatial patterns of both split-type and displacement-type SSW Z10 fields tend to reside in the lower-left quadrant with z 0 and z 1 larger than −2 (figures 6(a) and (b)).The highest chance of the vortex structure to occur is near z 0 = 0 and z 1 = −0.5 for splitting cases, corresponding to two vortices over North America and the North Atlantic, whereas that is near z 0 = −0.5 and z 1 = −0.5 for displacement cases, corresponding to a stretched and displaced vortex sitting over Europe.We also overlay the Z10 fields of split-type and displacement-type SSWs from MERRA-2 onto the VAE phase diagram (blue dots in figures 6(a) and (b)).Some observed vortices lie within the lower-left quadrant, where most of the simulated SSW vortices reside, while some others are located in the upper-right or upper-left quadrants, where fewer simulated SSW vortices appear.These results indicate that the vortex structure and its evolution in our WACCM6 simulations and those from reanalysis are somewhat different, although it is plausible that this is at least partly due to the small observational sample size rather than a model bias.
The same projections using WACCM6 SSWs are also made for the PCA phase diagram (figures 6(c) and (d)), which show less centered or organized regimes for both displacement and splitting types.The displacement-type events have a slightly larger chance of occurring in the lower-right quadrant but with a smaller percentage compared with the VAE outcome.The observed vortex structures during SSWs display a rather random distribution in the PCA phase diagram, unmatching the regime of higher occurrence formed by WACCM6 SSW events.The results suggest that the VAE is more capable of categorizing vortex structures into certain regimes in the phase diagram.To provide quantitative evidence that the VAE reconstructs better evolutions of SSWs than the PCA, we calculate the Rs and RMSEs as a function of time from day −10 to day 10 (figure 7) between the WACCM6 results and the reconstructions from PCA or the VAE model.As mentioned, the Rs quantifies the structural similarity between reconstructed Z10 fields and observed ones, whereas the RMSEs indicate the magnitude contrast between them.For both displacement and splitting events, the VAE gives higher Rs and lower RMSEs, averaged over 123 and 39 testing cases, respectively, than the PCA throughout most of the period, suggesting that the VAE performs better in terms of spatial reconstruction than the PCA.The only two exceptions are the RMSEs on day −3 for displacement events and on day 10 for splitting events, but the differences between them are rather small.The spread of VAE (light blue shading) is also smaller than that of PCA (light red shading), suggesting that a more statistically robust reconstruction is performed by the VAE.We notice that the Rs generally reduce with time in both VAE and PCA, while the RMSEs do not increase correspondingly.This means that the similarity between the reconstructed and observed Z10 patterns does not necessarily imply closer magnitude.
Finally, from the two cases we used to demonstrate the capability of VAE and PCA (figures 4 and 5), one may naively assume that a split-type SSW travels further in phase-space (i.e., greater overall variability during the event) than a displacement-type event.Therefore, to investigate if this is the case, we examine the total Euclidean distance D:

Discussion
Recent advances in deep learning have resulted in numerous innovative applications in climate and earth system sciences (e.g.The phase diagram also provides an innovative means to study the evolution of the vortex structure. The latent space of the VAE plays an essential role in this study.First, to enhance the physical interpretation of the VAE phase diagram, we try to connect the vortex transition in the diagram to the observed vortex changes.In the upper-left quadrant of the phase diagram, a rather circular vortex appears with its center at or near the North Pole.This structure manifests a low-amplitude wavenumber-1 disturbance or, geographiscally spealing, a stratospheric Aleutian anticyclone.In the lower-left corner of the diagram, the vortex becomes deepened, slightly elongated, and somewhat displaced toward Eurasia.In the upper-right corner and lower-right quadrant, the splitting structure of the polar vortex is illustrated accompanied by an amplified wavenumber-2 disturbance.The vortex in the lower-right quadrant turns elongated without clear displacement or splitting.However, according to the occurrence probability of SSW spatial patterns in the VAE phase diagram (figures 6(a) and (b)), the Z10 fields in the upper-right quadrant are absent in WACCM6 simulations and MERRA-2.These vortex structures are, thus, purely generated by the VAE.This indeed highlights the beauty of utilizing VAE to generate unseen or unknown vortex features and potentially shed new insights into the vortex structure and evolution.But a question immediately arises: do these vortex structures represent what has not, or has not yet, emerged in climate model simulations and in observations, or are they just data-driven artifacts?Looking into their spatial location, these elongated or split vortices are sitting above Alaska and eastern Siberia.Previous studies did not show these vortex patterns (e.g.Zhang et al 2016, Huang et al 2018), suggesting that they are unlikely to occur in the real world.Whether or not they will emerge in the future, in particular under a warmer climate and a corresponding colder stratosphere (e.g.Thompson et al 2012, Goessling and Bathiany 2016, Santer et al 2023), and how we exploit these generated vortices to enhance the understanding of the SSW evolution and underlying dynamics remain to be explored.
Next, we only use two means of the latent space to construct the phase diagram but do not consider the effect of exploiting the variances.From equation ( 2), the variance determines the magnitude and range of the Gaussian distribution to be fitted.Specifying different variance values could assess its influence on the vortex structure in the phase diagram and possibly give us further insights on the vortex structure transition.In addition, we only use two modes or dimensions of VAE to create the phase diagram.A previous study adopted 96 dimensions of VAE to best perform stratospheric polar vortex clustering (Krinitskiy et al 2019).Including more modes or dimensions could improve our VAE phase diagram, but would impede the intuitive illustration and physical interpretation would become more difficult.
Lastly, the VAE phase diagram provides another approach to manifest the discrepancy between SSWs in global climate model simulations and reanalysis data, which is suggested by figures 6 and 8.For both displacement-type and split-type SSWs, the simulated vortex structure occurs more frequently as the shapes displaced towards northern Europe but less frequently as the shapes of elongated vortices above Alaska and eastern Siberia, where a portion of observed vortex structures are located.Indeed, we find that the composite Z10 maps of all SSWs from the WACCM6 simulation show low geopotential heights over northern Europe.There may, therefore, exist model biases in the stratospheric polar vortex of WACCM6, which affects the latent space generated by VAE.In particular, biases in the structure of the stationary waves may be particularly important during the evolution of SSWs, when the stationary waves are amplified in the stratosphere.Further assessment on the probability distribution of WACCM6 and observed SSWs is needed to confirm this argument.While waiting for a sufficient number of observed SSWs to train a VAE seems unfeasible, training a VAE with SSWs from other state-of-the-art, high-top, large-ensemble climate models is more feasible and could be helpful in further understanding model biases.We would also like to use other reanalysis products with longer temporal coverage (and thus more SSWs) to test the VAE performance.
In addition, while we have investigated SSWs through the traditional lens of displacements and splits, the VAE phase diagram could provide a new framework through which to categorize the morphology of various types of SSWs.This may be particularly fruitful in terms of those SSWs for which distinguishing their type using a two-category system can be difficult.

Conclusion
In this study, we have successfully trained a VAE using the stratospheric polar vortex during SSWs from large-ensemble WACCM6 simulations and derived a phase diagram to characterize the spatiotemporal evolution of SSWs.Compared with the baseline results obtained from the PCA approach, the VAE gives rise to more organized vortex regimes in the phase diagram for both displacement-type and split-type SSW vortex structures than PCA.Moreover, the VAE can reconstruct a better spatial structure portraying vortex stretching, deformation, displacement, and splitting in the measures of spatial similarity and magnitude with respect to the observed vortices during most of the SSW evolution period.In particular, the VAE is more capable of reconstructing the splitting structure compared with the PCA reconstruction.In addition, the VAE phase diagram is able to capture differences between the model and observations, and thus offers a new platform to assess model biases.Based on the VAE framework, we anticipate future work to extend its applications to study the underlying dynamics of extreme stratospheric vortex phenomena and their change with climate change, as well as the associated two-way coupling to the troposphere that is crucial for improving and assessing subseasonal-to-seasonal prediction.

Figure 1 .
Figure 1.Illustration of the bilinear interpolation method applied to a Z10 field of a split-type SSW event from WACCM6 simulations.(a) The longitude-latitude map of the Northern Hemisphere in an orthographic projection with 90 • N at the center, with the region of color shadings indicating the domain to be interpolated.(b) The resultant 36 × 36 squared grid after using the bilinear interpolation method.The parallels on are respectively at 80, 60, and 40 degrees from the innermost to the outermost.

Figure 2 .
Figure2.A schematic of the VAE structure used in this study.The top row represents the combination of a decoder and an encoder, which form the VAE model.The middle row illustrates the details of the encoder, while the bottom row those of the decoder.The vertical solid line after each layer denotes the max-pooling/upsampling layers, which reduce or increase the dimension of pixeled images, whereas the cubes represent the convolutional layers with varying numbers of channels.The encoder reads in a 36 × 36 Z10 input image, compresses its information through the convolutional layers, and outputs two means (µ0, µ1) and two variances (v0, v1) that form the latent space.The decoder utilizes these sampled results based on Gaussian distribution and generates a 36 × 36 Z10 output through multiple convolutional layers.The mean squared error (MSE) between the original and generated images and the Kullback-Leibler divergence between the latent space and the Gaussian distribution are computed.These errors are backpropagated to train the encoder and decoder.The light blue color shading highlights the two means of the latent space, which is used to produce the phase diagram in this study.

Figure 3 .
Figure 3.This figure illustrates the corresponding Z10 field generated by the VAE decoder when given the varying values of z0 and z1 in the latent space.The green circle in each panel represents the circle of 60 • N.

Figure 4 .
Figure 4. (a) The evolution track of a displacement-type SSW event from WACCM6 simulations mapped onto the VAE phase diagram.(b) As in (a), but for the track mapped onto the PCA phase diagram.(c) The WACCM6 Z10 spatial patterns during days [−10, 10].(d) The VAE reconstructed Z10 fields during days [−10, 10].(e) As in (d), but for PCA reconstructed Z10 fields.The values in the subtitle parenthesis in (c) are the aspect ratio (r) and central latitude (ϕo) of this SSW event, whereas those in (c) and (d) denote the pattern correlation (R) and RMSE.The green circle in each panel of (c)-(e) represents the circle of 60 • N.

Figure 5 .
Figure 5.The same as in figure 4, but for a split-type SSW.

Figure 6 .
Figure 6.(a) The occurrence probability of the Z10 fields associated with all displacement-type SSWs from the WACCM6 testing dataset in the VAE phase diagram.The values are presented as percentages after normalized by the total number of events.(b) As in (a) but for split-type SSWs from WACCM6 simulations.Figures (c) and (d) are the same as (a) and (b), but for the PCA phase diagram.The blue dots are the Z10 fields from the reanalysis dataset mapped onto the VAE and PCA phase diagrams.

Figure 7 .
Figure 7.The pattern correlation (R) and root mean squared error (RMSE) for (a) and (b) the reconstructed displacement-type SSWs and (c) and (d) the reconstructed split-type SSWs from WACCM6 simulation against corresponding full Z10 during the period of days [−10, 10].The color shading indicates the range of one standard deviation.

Figure 8 .
Figure 8.(a) Distributions of distance traveled in the VAE phase diagram for all split-type (magenta bars) and displacement-type (blue bars) SSWs.(b) Distributions of distance traveled for all simulated SSWs (blue bars) and observational (magenta bars) SSWs.(c) and (d) The same as (a) and (b) but for the distributions of distance traveled in the PCA phase diagram.