tvec.py - pism - [fork] customized build of PISM, the parallel ice sheet model (tillflux branch)
(HTM) git clone git://src.adamsgaard.dk/pism
(DIR) Log
(DIR) Files
(DIR) Refs
(DIR) LICENSE
---
tvec.py (5666B)
---
1 # Copyright (C) 2011-2015, 2018 David Maxwell and Constantine Khroulev
2 #
3 # This file is part of PISM.
4 #
5 # PISM is free software; you can redistribute it and/or modify it under the
6 # terms of the GNU General Public License as published by the Free Software
7 # Foundation; either version 3 of the License, or (at your option) any later
8 # version.
9 #
10 # PISM is distributed in the hope that it will be useful, but WITHOUT ANY
11 # WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
12 # FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
13 # details.
14 #
15 # You should have received a copy of the GNU General Public License
16 # along with PISM; if not, write to the Free Software
17 # Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
18
19 """Functions and objects relating to interaction with :cpp:class:`IceModelVec`'s from python."""
20
21 import PISM
22
23
24 class Access(object):
25
26 """
27 Python context manager to simplify `IceModelVec` access and ghost communication.
28
29 In PISM C++ code, one uses :cpp:member:`IceModelVec::begin_access`/:cpp:member:`IceModelVec::end_access`
30 access pairs to delimit a code block allowing access to the contents of an :cpp:class`IcdeModelVec`.
31 If the contents of a ghosted vector were changed in the block, :cpp:member:`IceModelVec::update_ghosts` needs to
32 be called to synchronize the ghosts. Forgetting either an :cpp:member:`end_access` or an :cpp:member:`update_ghosts`
33 leads to bugs.
34
35 Python context managers are used in conjunction with ``with`` statements to execute code at the start
36 and end of a code block. A :class:`PISM.vec.Access` context manager is used to pair up
37 :cpp:member:`begin_access`/:cpp:member:`end_access` and to call :cpp:member:`update_ghosts` if needed:
38 Assuming that ``v1`` and ``v2`` are vectors::
39
40 grid = v1.get_grid()
41 with PISM.vec.Access(comm=v2, nocomm=v1):
42 for (i, j) in grid.points():
43 v2(i, j) = v1(i, j)**3
44
45 On entry into the ``with`` block, :cpp:member:`begin_access` is called for both ``v1`` and ``v2``.
46 On exit, :cpp:member:`end_access` is called for both ``v1`` and ``v2``, and :cpp:member:`update_ghosts`
47 is called for just ``v2``."""
48
49 def __init__(self, nocomm=None, comm=None):
50 """
51
52 :param nocomm: a vector or list of vectors to access such that
53 ghost communication *will not* occur when access is done.
54
55 :param comm: a vector or list of vectors to access such that
56 ghost communication *will* occur when access is done.
57
58 """
59 if not nocomm is None:
60 if isinstance(nocomm, list) or isinstance(nocomm, tuple):
61 self.nocomm = nocomm
62 else:
63 self.nocomm = [nocomm]
64 for v in self.nocomm:
65 v.begin_access()
66 else:
67 self.nocomm = None
68
69 if not comm is None:
70 if isinstance(comm, list) or isinstance(comm, tuple):
71 self.comm = comm
72 else:
73 self.comm = [comm]
74 for v in self.comm:
75 v.begin_access()
76 else:
77 self.comm = None
78
79 def __enter__(self):
80 pass
81
82 def __exit__(self, exc_type, exc_value, traceback):
83 if not self.nocomm is None:
84 for v in self.nocomm:
85 v.end_access()
86 self.nocomm = None
87
88 if not self.comm is None:
89 for v in self.comm:
90 v.end_access()
91 v.update_ghosts()
92 self.comm = None
93
94
95 def randVectorS(grid, scale, stencil_width=None):
96 """Create an :cpp:class:`IceModelVec2S` of normally distributed random entries.
97
98 :param grid: The :cpp:class:`IceGrid` to use for creating the vector.
99 :param scale: Standard deviation of normal distribution.
100 :param stencil_width: Ghost stencil width for the vector. Use ``None`` to indicate
101 an unghosted vector.
102
103 This function is not efficiently implemented.
104 """
105
106 if stencil_width is None:
107 stencil_width = 0
108 flag = PISM.WITHOUT_GHOSTS
109 else:
110 flag = PISM.WITH_GHOSTS
111
112 rv = PISM.IceModelVec2S(grid, 'rand vec', flag, stencil_width)
113
114 shape = (grid.xm(), grid.ym())
115 import numpy as np
116
117 r = np.random.normal(scale=scale, size=rv.shape())
118 with Access(nocomm=rv):
119 for (i, j) in grid.points():
120 rv[i, j] = r[j - grid.ys(), i - grid.xs()]
121
122 if stencil_width is not None:
123 rv.update_ghosts()
124
125 return rv
126
127
128 def randVectorV(grid, scale, stencil_width=None):
129 """Create an :cpp:class:`IceModelVec2V` of normally distributed random entries.
130
131 :param grid: The :cpp:class:`IceGrid` to use for creating the vector.
132 :param scale: Standard deviation of normal distribution.
133 :param stencil_width: Ghost stencil width for the vector. Use ``None`` to indicate
134 an unghosted vector.
135
136 This function is not efficiently implemented.
137 """
138
139 if stencil_width is None:
140 stencil_width = 0
141 flag = PISM.WITHOUT_GHOSTS
142 else:
143 flag = PISM.WITH_GHOSTS
144
145 rv = PISM.IceModelVec2V(grid, 'rand vec', flag, stencil_width)
146
147 import numpy as np
148
149 r_u = np.random.normal(scale=scale, size=rv.shape()[:-1])
150 r_v = np.random.normal(scale=scale, size=rv.shape()[:-1])
151 with Access(nocomm=rv):
152 for (i, j) in grid.points():
153 rv[i, j].u = r_u[j - grid.ys(), i - grid.xs()]
154 rv[i, j].v = r_v[j - grid.ys(), i - grid.xs()]
155
156 if stencil_width is not None:
157 rv.update_ghosts()
158
159 return rv