YetAnotherCoupler 3.1.1
Loading...
Searching...
No Matches
plotter.py

This component plots the configured fields on a regular grid using matplotlib.

1#!/usr/bin/env python3
2
3# Copyright (c) 2024 The YAC Authors
4#
5# SPDX-License-Identifier: BSD-3-Clause
6
7import os
8from yac import *
9import numpy as np
10import matplotlib.pyplot as plt
11
12class Plotter:
13 def __init__(self, variables, outdir = ".",
14 bounds = [0, 2*np.pi, -0.5*np.pi, 0.5*np.pi],
15 resolution = (360,180),
16 yac = None):
17 self.yac = yac or YAC.default_instance
18 self.comp = self.yac.predef_comp("plotter")
19 self.variables = variables
20 self.bounds = bounds
21 self.resolution = resolution
22 self.outdir = outdir
23
24 def setup(self):
25 self.x = np.linspace(self.bounds[0],self.bounds[1],self.resolution[0])
26 self.y = np.linspace(self.bounds[2],self.bounds[3],self.resolution[1])
27 grid = Reg2dGrid(f"plot_grid", self.x, self.y)
28 self.points = grid.def_points(Location.CORNER, self.x, self.y)
29
30 def def_couples(self):
31 nnn = InterpolationStack()
32 nnn.add_nnn(NNNReductionType.AVG, 1, 1.)
33
34 self.fields = []
35 for var in self.variables:
36 timestep = self.yac.get_field_timestep(*var)
37 collection_size = self.yac.get_field_collection_size(*var)
38 self.fields.append(Field.create(var[2], self.comp, self.points, 1,
39 timestep, TimeUnit.ISO_FORMAT))
40 self.yac.def_couple(*var,
41 "plotter", "plot_grid", var[2],
42 timestep, TimeUnit.ISO_FORMAT, 0, nnn)
43
44 for field in self.fields:
45 for i in range(field.collection_size):
46 os.makedirs(self.outdir +"/"+ f"{field.name}_{i}", exist_ok=True)
47
48 def step(self):
49 for field in self.fields:
50 time = field.datetime
51 print(f"plotting {field.name} at {time}")
52 buf, info = field.get()
53 for i in range(buf.shape[0]):
54 plt.imshow(buf[i,:].reshape((len(self.y),len(self.x)))[::-1,:],
55 extent=self.bounds)
56 plt.title(f"{field.name} - {time}")
57 plt.colorbar()
58 plt.savefig(self.outdir +"/"+ f"{field.name}_{i}/{time}.png")
59 plt.clf()
60 return field.datetime