[FEATURE] Add specific local domain for model forecasting

[ankitpatnala]

atmorep/atmorep/datasets/multifield_data_sampler.py

Lines 298 to 330 in 3efcb12

	def set_global( self, times, batch_size = None, token_overlap = [0, 0]) :
	''' generate patch/token positions for global grid '''
	token_overlap = np.array( token_overlap).astype(np.int64)
	# assumed that sanity checking that field data is consistent has been done
	ifield = 0
	field = self.fields[ifield]
	res = self.res
	side_len = np.array( [field[3][1] * field[4][1]*res[0], field[3][2] * field[4][2]*res[1]] )
	overlap = np.array([token_overlap[0]*field[4][1]*res[0],token_overlap[1]*field[4][2]*res[1]])
	side_len_2 = side_len / 2.
	assert all( overlap <= side_len_2), 'token_overlap too large for #tokens, reduce if possible'
	# generate tiles
	times_pos = []
	for ctime in times :
	lat = side_len_2[0].item()
	num_tiles_lat = 0
	while (lat + side_len_2[0].item()) < 180. :
	num_tiles_lat += 1
	lon = side_len_2[1].item() - overlap[1].item()/2.
	num_tiles_lon = 0
	while (lon - side_len_2[1]) < 360. :
	times_pos += [[*ctime, -lat + 90., np.mod(lon,360.) ]]
	lon += side_len[1].item() - overlap[1].item()
	num_tiles_lon += 1
	lat += side_len[0].item() - overlap[0].item()
	# add one additional row if no perfect tiling (sphere is toric in longitude so no special
	# handling necessary but not in latitude)
	# the added row is such that it goes exaclty down to the South pole and the offset North-wards

Is your feature request related to a problem? Please describe.
I want to see the pre-trained model's forecasting performance on a fixed domain (e.g CERRA). Currently, the function set_global() samples spatial lat and lon center points all across the globe.

Describe the solution you'd like
Add an argument domain with default parameter ([0,180],[0,360]). In case nothing is defined, it samples globally and when defined the start and end points will accordingly checked while sampling

Additional context

We need to focus on that specific location to see the model's performance on the domain before analyzing downscaling performance