1# -*- coding: utf-8 -*-
2
3from pyfr.mpiutil import get_comm_rank_root
4from pyfr.plugins.base import BasePlugin, init_csv
5
6
7class PseudoStatsPlugin(BasePlugin):
8    name = 'pseudostats'
9    systems = ['*']
10    formulations = ['dual']
11
12    def __init__(self, intg, cfgsect, prefix):
13        super().__init__(intg, cfgsect, prefix)
14
15        self.flushsteps = self.cfg.getint(self.cfgsect, 'flushsteps', 500)
16
17        self.count = 0
18        self.stats = []
19        self.tprev = intg.tcurr
20
21        fvars = ','.join(intg.system.elementscls.convarmap[self.ndims])
22
23        # MPI info
24        comm, rank, root = get_comm_rank_root()
25
26        # The root rank needs to open the output file
27        if rank == root:
28            self.outf = init_csv(self.cfg, cfgsect, 'n,t,i,' + fvars)
29        else:
30            self.outf = None
31
32    def __call__(self, intg):
33        # Process the sequence of pseudo-residuals
34        for (npiter, iternr, resid) in intg.pseudostepinfo:
35            resid = resid or ('-',)*intg.system.nvars
36            self.stats.append((npiter, self.tprev, iternr) + resid)
37
38        # Update the total step count and save the current time
39        self.count += len(intg.pseudostepinfo)
40        self.tprev = intg.tcurr
41
42        # If we're the root rank then output
43        if self.outf:
44            for s in self.stats:
45                print(*s, sep=',', file=self.outf)
46
47            # Periodically flush to disk
48            if intg.nacptsteps % self.flushsteps == 0:
49                self.outf.flush()
50
51        # Reset the stats
52        self.stats = []
53