Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/dataset/topo/dataset.py: 89%

1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

4# ---|P------/S----------~Lg---------- 

5 

6''' 

7Functionality to handle tiled global topography datasets. 

8''' 

9 

10import math 

11import logging 

12import os.path as op 

13 

14import numpy as num 

15 

16from . import tile 

17from ..util import get_download_callback 

18from pyrocko import util 

19 

20try: 

21 range = xrange 

22except NameError: 

23 pass 

24 

25logger = logging.getLogger('pyrocko.dataset.topo.dataset') 

26 

27 

28class TiledGlobalDataset(object): 

29 ''' 

30 Tiled global dataset. 

31 ''' 

32 

33 def __init__(self, name, nx, ny, ntx, nty, dtype, data_dir=None, 

34 citation=None, region=None): 

35 

36 # dataset geometry (including overlap/endpoints) 

37 self.nx = int(nx) 

38 self.ny = int(ny) 

39 self.xmin = -180. 

40 self.xmax = 180. 

41 self.ymin = -90. 

42 self.ymax = 90. 

43 self.dx = (self.xmax-self.xmin) / (self.nx - 1) 

44 self.dy = (self.ymax-self.ymin) / (self.ny - 1) 

45 

46 # tile geometry (including overlap) 

47 self.ntx = int(ntx) 

48 self.nty = int(nty) 

49 self.stx = (self.ntx - 1) * self.dx 

50 self.sty = (self.nty - 1) * self.dy 

51 

52 self.ntilesx = (self.nx - 1) // (self.ntx - 1) 

53 self.ntilesy = (self.ny - 1) // (self.nty - 1) 

54 

55 self.name = name 

56 self.dtype = dtype 

57 self.data_dir = data_dir 

58 self.citation = citation 

59 if region is not None: 

60 self.region = self.positive_region(region) 

61 else: 

62 self.region = None 

63 

64 def covers(self, region): 

65 if self.region is None: 

66 return True 

67 

68 a = self.region 

69 b = self.positive_region(region) 

70 return (( 

71 (a[0] <= b[0] and b[1] <= b[1]) or 

72 (a[0] <= b[0]+360. and b[1]+360 <= a[1])) 

73 and a[2] <= b[2] and b[3] <= a[3]) 

74 

75 def is_suitable(self, region, dmin, dmax): 

76 d = 360. / (self.nx - 1) 

77 return self.covers(region) and dmin <= d <= dmax 

78 

79 def download_file(self, url, fpath, username=None, password=None): 

80 # TODO: Add logstring 

81 util.download_file( 

82 url, fpath, username, password, 

83 status_callback=get_download_callback( 

84 'Downloading %s topography...' % self.__class__.__name__)) 

85 

86 def x(self): 

87 return self.xmin + num.arange(self.nx) * self.dx 

88 

89 def y(self): 

90 return self.ymin + num.arange(self.ny) * self.dy 

91 

92 def positive_region(self, region): 

93 xmin, xmax, ymin, ymax = [float(x) for x in region] 

94 

95 assert -180. - 360. <= xmin < 180. 

96 assert -180. < xmax <= 180. + 360. 

97 assert -90. <= ymin < 90. 

98 assert -90. < ymax <= 90. 

99 

100 if xmax < xmin: 

101 xmax += 360. 

102 

103 if xmin < -180.: 

104 xmin += 360. 

105 xmax += 360. 

106 

107 return (xmin, xmax, ymin, ymax) 

108 

109 def tile_indices(self, region): 

110 xmin, xmax, ymin, ymax = self.positive_region(region) 

111 itxmin = int(math.floor((xmin - self.xmin) / self.stx)) 

112 itxmax = int(math.ceil((xmax - self.xmin) / self.stx)) 

113 if itxmin == itxmax: 

114 itxmax += 1 

115 itymin = int(math.floor((ymin - self.ymin) / self.sty)) 

116 itymax = int(math.ceil((ymax - self.ymin) / self.sty)) 

117 if itymin == itymax: 

118 if itymax != self.ntilesy: 

119 itymax += 1 

120 else: 

121 itymin -= 1 

122 

123 indices = [] 

124 for ity in range(itymin, itymax): 

125 for itx in range(itxmin, itxmax): 

126 indices.append((itx % self.ntilesx, ity)) 

127 

128 return indices 

129 

130 def get_tile(self, itx, ity): 

131 return None 

132 

133 def get(self, region): 

134 if len(region) == 2: 

135 x, y = region 

136 t = self.get((x, x, y, y)) 

137 if t is not None: 

138 return t.get(x, y) 

139 else: 

140 return None 

141 

142 indices = self.tile_indices(region) 

143 tiles = [] 

144 for itx, ity in indices: 

145 t = self.get_tile(itx, ity) 

146 if t: 

147 tiles.append(t) 

148 

149 return tile.combine(tiles, region) 

150 

151 def get_with_repeat(self, region): 

152 xmin, xmax, ymin, ymax = region 

153 ymin2 = max(-90., ymin) 

154 ymax2 = min(90., ymax) 

155 region2 = xmin, xmax, ymin2, ymax2 

156 t = self.get(region2) 

157 if t is not None and region2 != region: 

158 t.yextend_with_repeat(ymin, ymax) 

159 

160 return t 

161 

162 

163class DecimatedTiledGlobalDataset(TiledGlobalDataset): 

164 ''' 

165 Decimated variant of a tiled global dataset. 

166 ''' 

167 

168 def __init__(self, name, base, ndeci, data_dir=None, ntx=None, nty=None): 

169 

170 assert ndeci % 2 == 0 

171 assert (base.nx - 1) % ndeci == 0 

172 assert (base.ny - 1) % ndeci == 0 

173 

174 nx = (base.nx - 1) // ndeci + 1 

175 ny = (base.ny - 1) // ndeci + 1 

176 

177 if ntx is None: 

178 ntx = base.ntx 

179 

180 if nty is None: 

181 nty = base.nty 

182 

183 assert (nx - 1) % (ntx - 1) == 0 

184 assert (ny - 1) % (nty - 1) == 0 

185 

186 if data_dir is None: 

187 data_dir = op.join(base.data_dir, 'decimated_%i' % ndeci) 

188 

189 TiledGlobalDataset.__init__(self, name, nx, ny, ntx, nty, base.dtype, 

190 data_dir=data_dir, citation=base.citation, 

191 region=base.region) 

192 

193 self.base = base 

194 self.ndeci = ndeci 

195 

196 def make_tile(self, itx, ity, fpath): 

197 nh = self.ndeci // 2 

198 xmin = self.xmin + itx*self.stx - self.base.dx * nh 

199 xmax = self.xmin + (itx+1)*self.stx + self.base.dx * nh 

200 ymin = self.ymin + ity*self.sty - self.base.dy * nh 

201 ymax = self.ymin + (ity+1)*self.sty + self.base.dy * nh 

202 

203 t = self.base.get_with_repeat((xmin, xmax, ymin, ymax)) 

204 if t is not None: 

205 t.decimate(self.ndeci) 

206 

207 util.ensuredirs(fpath) 

208 with open(fpath, 'w') as f: 

209 if t is not None: 

210 t.data.tofile(f) 

211 

212 def make_if_needed(self, itx, ity): 

213 assert 0 <= itx < self.ntilesx 

214 assert 0 <= ity < self.ntilesy 

215 

216 fn = '%02i.%02i.bin' % (ity, itx) 

217 fpath = op.join(self.data_dir, fn) 

218 if not op.exists(fpath): 

219 logger.info('making decimated tile: %s (%s)' % (fn, self.name)) 

220 self.make_tile(itx, ity, fpath) 

221 

222 def get_tile(self, itx, ity): 

223 assert 0 <= itx < self.ntilesx 

224 assert 0 <= ity < self.ntilesy 

225 

226 self.make_if_needed(itx, ity) 

227 

228 fn = '%02i.%02i.bin' % (ity, itx) 

229 fpath = op.join(self.data_dir, fn) 

230 with open(fpath, 'r') as f: 

231 data = num.fromfile(f, dtype=self.dtype) 

232 

233 if data.size == 0: 

234 return None 

235 

236 assert data.size == self.ntx*self.nty 

237 data = data.reshape((self.ntx, self.nty)) 

238 

239 return tile.Tile( 

240 self.xmin + itx*self.stx, 

241 self.ymin + ity*self.sty, 

242 self.dx, self.dx, data) 

243 

244 def make_all_missing(self): 

245 for ity in range(self.ntilesy): 

246 for itx in range(self.ntilesx): 

247 self.make_if_needed(itx, ity)