datagenerator.Horizons

   1import os
   2import math
   3import numpy as np
   4from scipy import stats
   5from scipy.interpolate import griddata
   6from datagenerator.Parameters import Parameters
   7from itertools import groupby
   8
   9class Horizons:
  10    def __init__(self, parameters):
  11        self.cfg = parameters
  12        self.max_layers = 0
  13
  14    def insert_feature_into_horizon_stack(self, feature, layer_number, maps):
  15        """
  16        Insert an object into horizons which is to be inserted into a single layer
  17
  18        Made for fans, but should work for any object.
  19        Shallower layer thicknesses are adjusted to insert the feature (clipped downwards onto the feature).
  20        """
  21        layer_count = 1
  22        if feature.ndim == 3:
  23            layer_count = feature.shape[-1]
  24        new_maps = np.zeros(
  25            (
  26                self.cfg.cube_shape[0],
  27                self.cfg.cube_shape[1],
  28                maps.shape[2] + layer_count,
  29            )
  30        )
  31        # Shift horizons below layer downwards by 1 to insert object in specified layer
  32        new_maps[..., layer_number + layer_count :] = maps[..., layer_number:]
  33        # Insert object into layer_number
  34        new_maps[..., layer_number] = maps[..., layer_number] - feature
  35        # copy horizons above layer_number into the new set of maps
  36        new_maps[..., :layer_number] = maps[..., :layer_number]
  37
  38        # Don't allow negative thickness after inserting the fan
  39        for i in range(new_maps.shape[-1] - 1, 1, -1):
  40            layer_thickness = new_maps[..., i] - new_maps[..., i - 1]
  41            if np.min(layer_thickness) < 0:
  42                np.clip(layer_thickness, 0, a_max=None, out=layer_thickness)
  43                new_maps[..., i - 1] = new_maps[..., i] - layer_thickness
  44        # Increase the max number of layers in model by 1
  45        self.max_layers += layer_count
  46        return new_maps
  47
  48    def insert_seafloor(self, maps):
  49        if self.cfg.verbose:
  50            print("\n ... inserting 'water bottom' reflector in work cube ...\n")
  51        wb_time_map = maps[:, :, 1] - 1.5
  52        wb_stats = [
  53            f"{x * self.cfg.digi / self.cfg.infill_factor:.2f}"
  54            for x in [wb_time_map.min(), wb_time_map.mean(), wb_time_map.max()]
  55        ]
  56        self.cfg.write_to_logfile("Seabed Min: {}, Mean: {}, Max: {}".format(*wb_stats))
  57
  58        self.cfg.write_to_logfile(
  59            msg=None,
  60            mainkey="model_parameters",
  61            subkey="seabed_min",
  62            val=wb_time_map.min(),
  63        )
  64        self.cfg.write_to_logfile(
  65            msg=None,
  66            mainkey="model_parameters",
  67            subkey="seabed_mean",
  68            val=wb_time_map.mean(),
  69        )
  70        self.cfg.write_to_logfile(
  71            msg=None,
  72            mainkey="model_parameters",
  73            subkey="seabed_max",
  74            val=wb_time_map.max(),
  75        )
  76
  77        maps[:, :, 0] = wb_time_map.copy()
  78        return maps
  79
  80    def create_random_net_over_gross_map(
  81        self, avg=(0.45, 0.9), stdev=(0.01, 0.05), octave=9
  82    ):
  83        random_net_over_gross_map = self._perlin(base=None, octave=octave)
  84
  85        avg_net_over_gross = np.random.uniform(*avg)
  86        avg_net_over_gross_stdev = np.random.uniform(*stdev)
  87
  88        # set stdev and mean of map to desired values
  89        random_net_over_gross_map -= random_net_over_gross_map.mean()
  90        random_net_over_gross_map *= (
  91            avg_net_over_gross_stdev / random_net_over_gross_map.std()
  92        )
  93        random_net_over_gross_map += avg_net_over_gross
  94        # clip to stay within a reasonable N/G range
  95        random_net_over_gross_map = random_net_over_gross_map.clip(*avg)
  96        return random_net_over_gross_map
  97
  98    def _perlin(self, base=None, octave=1, lac=1.9, do_rotate=True):
  99        import noise
 100
 101        xsize = self.cfg.cube_shape[0]
 102        ysize = self.cfg.cube_shape[1]
 103        if base is None:
 104            base = np.random.randint(255)
 105        # randomly rotate image
 106        if do_rotate:
 107            number_90_deg_rotations = 0
 108            fliplr = False
 109            flipud = False
 110            if xsize == ysize and np.random.binomial(1, 0.5) == 1:
 111                number_90_deg_rotations = int(np.random.uniform(1, 4))
 112                # temp = np.rot90(temp, number_90_deg_rotations)
 113            # randomly flip left and right, top and bottom
 114            if np.random.binomial(1, 0.5) == 1:
 115                fliplr = True
 116            if np.random.binomial(1, 0.5) == 1:
 117                flipud = True
 118
 119        temp = np.array(
 120            [
 121                [
 122                    noise.pnoise2(
 123                        float(i) / xsize,
 124                        float(j) / ysize,
 125                        lacunarity=lac,
 126                        octaves=octave,
 127                        base=base,
 128                    )
 129                    for j in range(ysize)
 130                ]
 131                for i in range(xsize)
 132            ]
 133        )
 134        temp = np.rot90(temp, number_90_deg_rotations)
 135        if fliplr:
 136            temp = np.fliplr(temp)
 137        if flipud:
 138            temp = np.flipud(temp)
 139        return temp
 140
 141    def _fit_plane_strike_dip(self, azimuth, dip, grid_shape, verbose=False):
 142        # Fits a plane given dip and max dip direction (azimuth)
 143        # - create a point at the center of the grid, elevation is zero
 144        xyz1 = np.array([grid_shape[0] / 2.0, grid_shape[1] / 2.0, 0.0])
 145        # y = np.array([0., 1., 1.])
 146        # - create a point in the strike direction at same elevation of zero
 147        strike_angle = azimuth + 90.0
 148        if strike_angle > 360.0:
 149            strike_angle -= 360.0
 150        if strike_angle > 180.0:
 151            strike_angle -= 180.0
 152        strike_angle *= np.pi / 180.0
 153        distance = np.min(grid_shape) / 4.0
 154        x = distance * math.cos(strike_angle) + grid_shape[0] / 2.0
 155        y = distance * math.sin(strike_angle) + grid_shape[1] / 2.0
 156        xyz2 = np.array([x, y, 0.0])
 157        # - create a point in the max dip direction
 158        dip_angle = dip * 1.0
 159        if dip_angle > 360.0:
 160            dip_angle -= 360.0
 161        if dip_angle > 180.0:
 162            dip_angle -= 180.0
 163        dip_angle *= np.pi / 180.0
 164        strike_angle = azimuth
 165        if strike_angle > 360.0:
 166            strike_angle -= 360.0
 167        if strike_angle > 180.0:
 168            strike_angle -= 180.0
 169            dip_elev = -distance * math.sin(dip_angle) * math.sqrt(2.0)
 170        else:
 171            dip_elev = distance * math.sin(dip_angle) * math.sqrt(2.0)
 172        strike_angle *= np.pi / 180.0
 173        x = distance * math.cos(strike_angle) + grid_shape[0] / 2.0
 174        y = distance * math.sin(strike_angle) + grid_shape[1] / 2.0
 175        xyz3 = np.array([x, y, dip_elev])
 176        # - combine 3 points into single array
 177        xyz = np.vstack((xyz1, xyz2, xyz3))
 178        # - fit three points to a plane and compute elevation for all grids on surface
 179        a, b, c = self.fit_plane_lsq(xyz)
 180        z = self.eval_plane(range(grid_shape[0]), range(grid_shape[1]), a, b, c)
 181        # - make plot if not quiet
 182        if verbose:
 183            import os
 184            from datagenerator.util import import_matplotlib
 185
 186            plt = import_matplotlib()
 187            print(f"strike angle = {strike_angle}")
 188            print(f"dip angle = {dip_angle}")
 189            print(f"points are: {xyz}")
 190            plt.figure(1)
 191            plt.clf()
 192            plt.grid()
 193            plt.imshow(z, origin="origin")
 194            plt.plot(xyz[:, 0], xyz[:, 1], "yo")
 195            plt.colorbar()
 196            plt.savefig(os.path.join(self.cfg.work_subfolder, "dipping_plane.png"))
 197            plt.close()
 198        return z
 199
 200    @staticmethod
 201    def fit_plane_lsq(xyz):
 202        # Fits a plane to a point cloud,
 203        # Where Z = aX + bY + c        ----Eqn #1
 204        # Rearranging Eqn1: aX + bY -Z +c =0
 205        # Gives normal (a,b,-1)
 206        # Normal = (a,b,-1)
 207        # [rows, cols] = xyz.shape
 208        rows = xyz.shape[0]
 209        g = np.ones((rows, 3))
 210        g[:, 0] = xyz[:, 0]  # X
 211        g[:, 1] = xyz[:, 1]  # Y
 212        z = xyz[:, 2]
 213        (a, b, c), _, _, _ = np.linalg.lstsq(g, z, rcond=-1)
 214        normal = np.array([a, b, c])
 215        return normal
 216
 217    @staticmethod
 218    def eval_plane(x, y, a, b, c):
 219        # evaluates and returns Z for each input X,Y
 220        z = np.zeros((len(x), len(y)), "float")
 221        for i in range(len(x)):
 222            for j in range(len(y)):
 223                z[i, j] = a * x[i] + b * y[j] + c
 224        return z
 225
 226    @staticmethod
 227    def halton(dim, nbpts):
 228        import math
 229
 230        h = np.empty(nbpts * dim)
 231        h.fill(np.nan)
 232        p = np.empty(nbpts)
 233        p.fill(np.nan)
 234        p1 = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31]
 235        lognbpts = math.log(nbpts + 1)
 236        for i in range(dim):
 237            b = p1[i]
 238            n = int(math.ceil(lognbpts / math.log(b)))
 239            for t in range(n):
 240                p[t] = pow(b, -(t + 1))
 241
 242            for j in range(nbpts):
 243                d = j + 1
 244                sum_ = math.fmod(d, b) * p[0]
 245                for t in range(1, n):
 246                    d = math.floor(d / b)
 247                    sum_ += math.fmod(d, b) * p[t]
 248
 249                h[j * dim + i] = sum_
 250
 251        return h.reshape(nbpts, dim)
 252
 253    @staticmethod
 254    def rotate_point(x, y, angle_in_degrees):
 255        """Calculate new coordinates for point after coordinate rotation about (0,0) by angleDegrees"""
 256        angle = angle_in_degrees * np.pi / 180.0
 257        x1 = math.cos(angle) * x + math.sin(angle) * y
 258        y1 = -math.sin(angle) * x + math.cos(angle) * y
 259        return x1, y1
 260
 261    def convert_map_from_samples_to_units(self, maps):
 262        """Use digi to convert a copy of the provided maps from samples to units"""
 263        converted_maps = maps.copy()
 264        converted_maps *= float(self.cfg.digi)
 265        return converted_maps
 266
 267    def write_maps_to_disk(self, horizons, name):
 268        """Write horizons to disk."""
 269        fname = os.path.join(self.cfg.work_subfolder, name)
 270        np.save(fname, horizons)
 271
 272    def write_onlap_episodes(
 273        self, onlap_horizon_list, depth_maps_gaps, depth_maps_infilled, n=35
 274    ):
 275        """Write gapped horizons with onlaps + n shallower horizons to a separate files."""
 276        tilting_zmaps = list()
 277        interval_thickness_size = depth_maps_gaps.shape[0] * depth_maps_gaps.shape[1]
 278        for ihor in onlap_horizon_list:
 279            for jhor in range(ihor - n, ihor + 1):
 280                if jhor > 0:
 281                    # compute thickness compared to next horizon using infilled depth maps
 282                    interval_thickness = (
 283                        depth_maps_infilled[:, :, jhor + 1]
 284                        - depth_maps_infilled[:, :, jhor]
 285                    )
 286                    # compute percentage of interval filled with zeros
 287                    interval_thickness_zeros_size = interval_thickness[
 288                        interval_thickness < 1.0e-5
 289                    ].shape[0]
 290                    pct_zeros = (
 291                        float(interval_thickness_zeros_size) / interval_thickness_size
 292                    )
 293                    if pct_zeros > 0.02:
 294                        tilting_zmaps.append(jhor)
 295        # Make the tilting maps unique and sorted
 296        tilting_zmaps = list(set(tilting_zmaps))
 297        tilting_zmaps.sort()
 298        onlaps = np.zeros(
 299            (depth_maps_gaps.shape[0], depth_maps_gaps.shape[1], len(tilting_zmaps))
 300        )
 301        for count, h in enumerate(tilting_zmaps):
 302            onlaps[..., count] = depth_maps_gaps[..., h]
 303
 304        # Write to disk
 305        self.write_maps_to_disk(onlaps * self.cfg.digi, "depth_maps_onlaps")
 306
 307        if self.cfg.hdf_store:
 308            # Write onlap maps to hdf
 309            from datagenerator.util import write_data_to_hdf
 310
 311            write_data_to_hdf(
 312                "depth_maps_onlaps", onlaps * self.cfg.digi, self.cfg.hdf_master
 313            )
 314
 315    def write_fan_horizons(self, fan_horizon_list, depth_maps):
 316        """Write fan layers to a separate file."""
 317        z = depth_maps.copy()
 318        fan_horizons = np.zeros((z.shape[0], z.shape[1], len(fan_horizon_list)))
 319        for count, horizon in enumerate(fan_horizon_list):
 320            thickness_map = z[..., horizon + 1] - z[..., horizon]
 321            z[..., horizon][thickness_map <= 0] = np.nan
 322            if self.cfg.verbose:
 323                print(
 324                    f"Fan horizon number: {horizon} Max thickness: {np.nanmax(thickness_map)}"
 325                )
 326            fan_horizons[..., count] = z[..., horizon]
 327        self.write_maps_to_disk(fan_horizons, "depth_maps_fans")
 328
 329
 330class RandomHorizonStack(Horizons):
 331    def __init__(self, parameters):
 332        self.cfg = parameters
 333        self.depth_maps = None
 334        self.depth_maps_gaps = None
 335        self.max_layers = 0
 336
 337        # Look up tables
 338        self.thicknesses = None
 339        self.onlaps = None
 340        self.channels = None
 341        self.dips = None
 342        self.azimuths = None
 343        self.facies = None
 344
 345    def _generate_lookup_tables(self):
 346        # Thicknesses
 347        self.thicknesses = stats.gamma.rvs(4.0, 2, size=self.cfg.num_lyr_lut)
 348        # Onlaps
 349        onlap_layer_list = np.sort(
 350            np.random.uniform(
 351                low=5, high=200, size=int(np.random.triangular(1, 4, 7) + 0.5)
 352            ).astype("int")
 353        )
 354        # Dips
 355        self.dips = (
 356            (1.0 - random.power(100, self.cfg.num_lyr_lut))
 357            * 7.0
 358            * self.cfg.dip_factor_max
 359        )
 360        # Azimuths
 361        self.azimuths = np.random.uniform(
 362            low=0.0, high=360.0, size=(self.cfg.num_lyr_lut,)
 363        )
 364        # Channels
 365        self.channels = np.random.binomial(1, 3.0 / 100.0, self.cfg.num_lyr_lut)
 366
 367        if self.cfg.verbose:
 368            print("self.cfg.num_lyr_lut = ", self.cfg.num_lyr_lut)
 369            print("onlap_layer_list = ", onlap_layer_list)
 370        onlap_array_dim = int(500 / 1250 * self.cfg.cube_shape[2])
 371        self.onlaps = np.zeros(onlap_array_dim, "int")
 372        self.onlaps[onlap_layer_list] = 1
 373
 374        if not self.cfg.include_channels:
 375            self.channels *= 0
 376        else:
 377            # Make sure channel episodes don't overlap too much
 378            for i in range(len(self.channels)):
 379                if self.channels[i] == 1:
 380                    self.channels[i + 1 : i + 6] *= 0
 381
 382        if self.cfg.verbose:
 383            print(
 384                f"Number of onlapping flags: {self.onlaps[self.onlaps == 1].shape[0]}"
 385            )
 386            print(
 387                f" ... horizon number for first onlap episode = {np.argmax(self.onlaps)}"
 388            )
 389            print(
 390                f" ... number of channelFlags: {self.channels[:250][self.channels[:250] == 1].shape[0]}"
 391            )
 392            print(
 393                f" ... horizon number for first channel episode: {np.argmax(self.channels)}"
 394            )
 395
 396    def _random_layer_thickness(self):
 397        rand_oct = int(np.random.triangular(left=1.3, mode=2.65, right=5.25))
 398
 399        low_thickness_factor = np.random.triangular(left=0.05, mode=0.2, right=0.95)
 400
 401        high_thickness_factor = np.random.triangular(left=1.05, mode=1.8, right=2.2)
 402
 403        thickness_factor_map = self._perlin(octave=rand_oct)
 404        thickness_factor_map -= thickness_factor_map.mean()
 405        thickness_factor_map *= (high_thickness_factor - low_thickness_factor) / (
 406            thickness_factor_map.max() - thickness_factor_map.min()
 407        )
 408        thickness_factor_map += 1.0
 409        if thickness_factor_map.min() < 0.0:
 410            thickness_factor_map -= thickness_factor_map.min() - 0.05
 411
 412        if self.cfg.verbose:
 413            print(
 414                f" {thickness_factor_map.mean()}, "
 415                f"{thickness_factor_map.min()}, "
 416                f" {thickness_factor_map.max()}, "
 417                f" {thickness_factor_map.std()}"
 418            )
 419
 420        return thickness_factor_map
 421
 422    def _generate_random_depth_structure_map(
 423        self,
 424        dip_range=(0.25, 1.0),
 425        num_points_range=(2, 25),
 426        elevation_std=100.0,
 427        zero_at_corners=True,
 428        initial=False,
 429    ):
 430        ############################################################################
 431        # generate a 2D array representing a depth (or TWT) structure map.
 432        # - grid_size controls output size in (x,y), units are samples
 433        # - dip_range controls slope of dipping plane upon which random residual points are placed
 434        # - num_points controls number of randomly positioned points are used for residual grid
 435        # - elevation_std controls std dev for residual elevation of random points in residual grid
 436        ############################################################################
 437        grid_size = self.cfg.cube_shape[:2]
 438        # define output grid (padded by 15% in each direction)
 439        xi = np.linspace(-0.15, 1.15, int(grid_size[0] * 1.3))
 440        yi = np.linspace(-0.15, 1.15, int(grid_size[1] * 1.3))
 441
 442        # start with a gently dipping plane similar to that found on passive shelf margins (< 1 degree dip)
 443        azimuth = random.uniform(0.0, 360.0)
 444        dip = random.uniform(dip_range[0], dip_range[1])
 445
 446        # build a residual surface to add to the dipping plane
 447        number_halton_points = int(random.uniform(100, 500) + 0.5)
 448        number_random_points = int(
 449            random.uniform(num_points_range[0], num_points_range[1]) + 0.5
 450        )
 451        z = np.random.rand(number_random_points)
 452        z -= z.mean()
 453        if initial:
 454            elevation_std = self.cfg.initial_layer_stdev
 455        z *= elevation_std / z.std()
 456
 457        dipping_plane = self._fit_plane_strike_dip(
 458            azimuth, dip, grid_size, verbose=False
 459        )
 460        xx = self.halton(2, number_halton_points)[-number_random_points:]
 461
 462        # make up some randomly distributed data
 463        # - adjust for padding
 464        xx *= xi[-1] - xi[0]
 465        xx += xi[0]
 466        x = xx[:, 0]
 467        y = xx[:, 1]
 468
 469        if zero_at_corners:
 470            x = np.hstack((x, [-0.15, -0.15, 1.15, 1.15]))
 471            y = np.hstack((y, [-0.15, 1.15, -0.15, 1.15]))
 472            z = np.hstack((z, [-0.15, -0.15, -0.15, -0.15]))
 473
 474        # grid the data.
 475        zi = griddata(
 476            np.column_stack((x, y)),
 477            z,
 478            (xi[:, np.newaxis], yi[np.newaxis, :]),
 479            method="cubic",
 480        )
 481        zi = zi.reshape(xi.shape[0], yi.shape[0])
 482        xi_min_index = np.argmin((xi - 0.0) ** 2)
 483        xi_max_index = xi_min_index + grid_size[0]
 484        yi_min_index = np.argmin((yi - 0.0) ** 2)
 485        yi_max_index = yi_min_index + grid_size[1]
 486        zi = zi[xi_min_index:xi_max_index, yi_min_index:yi_max_index]
 487        # add to the gently sloping plane
 488        zi += dipping_plane
 489
 490        if initial:
 491            zi += dipping_plane - dipping_plane.min()
 492            zi += self.cfg.cube_shape[-1] * self.cfg.infill_factor - zi.min()
 493            zi_argmin_i, zi_argmin_j = np.unravel_index(zi.argmin(), zi.shape)
 494            if self.cfg.verbose:
 495                print(
 496                    f"\tIndices for shallowest point in cube: {zi_argmin_i}, {zi_argmin_j}"
 497                )
 498            self.cfg.write_to_logfile(
 499                f"number_random_points: {number_random_points:.0f}"
 500            )
 501            self.cfg.write_to_logfile(f"dip angle: {dip:.2f}")
 502            self.cfg.write_to_logfile(
 503                f"dipping_plane min: {dipping_plane.min():.2f}, mean: {dipping_plane.mean():.2f},"
 504                f" max: {dipping_plane.max():.2f}"
 505            )
 506            self.cfg.write_to_logfile(
 507                f"zi min: {zi.min():.2f}, mean: {zi.mean():.2f}, max: {zi.max():.2f}"
 508            )
 509
 510            self.cfg.write_to_logfile(
 511                msg=None,
 512                mainkey="model_parameters",
 513                subkey="number_random_points",
 514                val=number_random_points,
 515            )
 516            self.cfg.write_to_logfile(
 517                msg=None,
 518                mainkey="model_parameters",
 519                subkey="dip_angle",
 520                val=number_random_points,
 521            )
 522            self.cfg.write_to_logfile(
 523                msg=None,
 524                mainkey="model_parameters",
 525                subkey="dipping_plane_min",
 526                val=dipping_plane.min(),
 527            )
 528            self.cfg.write_to_logfile(
 529                msg=None,
 530                mainkey="model_parameters",
 531                subkey="dipping_plane_mean",
 532                val=dipping_plane.mean(),
 533            )
 534            self.cfg.write_to_logfile(
 535                msg=None,
 536                mainkey="model_parameters",
 537                subkey="dipping_plane_max",
 538                val=dipping_plane.max(),
 539            )
 540            self.cfg.write_to_logfile(
 541                msg=None, mainkey="model_parameters", subkey="zi_min", val=zi.min()
 542            )
 543            self.cfg.write_to_logfile(
 544                msg=None, mainkey="model_parameters", subkey="zi_mean", val=zi.mean()
 545            )
 546            self.cfg.write_to_logfile(
 547                msg=None, mainkey="model_parameters", subkey="zi_max", val=zi.max()
 548            )
 549
 550        return xx, zi
 551
 552    def _create_thickness_map(self, random_thickness_factor_map, layer_number):
 553        """
 554        Create a random_thickness_map for a given layer.
 555
 556        Creates a dipping_plane with differential dip from previous layer and computes a new thickness_map which is
 557         always positive (i.e. no erosion allowed)
 558
 559        Parameters
 560        ----------
 561        random_thickness_factor_map : ndarray - a random thickness factor map for this layer
 562        layer_number : int - the layer number being built in the stack of horizons
 563
 564        Returns
 565        -------
 566        thickness_map : ndarray - the thickness_map for the given layer
 567        """
 568        dipping_plane = self._fit_plane_strike_dip(
 569            azimuth=self.azimuths[layer_number],
 570            dip=self.dips[layer_number],
 571            grid_shape=self.cfg.cube_shape[:2],
 572            verbose=False,
 573        )
 574        dipping_plane -= dipping_plane.min()
 575        if self.cfg.verbose:
 576            print(
 577                f"azi, dip, dipping_plane min/mean/max = {self.azimuths[layer_number]}, {self.dips[layer_number]}, "
 578                f"{dipping_plane.min():.2f}, {dipping_plane.mean():.2f}, {dipping_plane.max():.2f}"
 579            )
 580
 581        # compute new thickness map. For now, don't allow erosion (i.e. thickness_map is always positive)
 582        thickness_map = (
 583            (
 584                np.ones(self.cfg.cube_shape[:2], "float")
 585                * self.thicknesses[layer_number]
 586                - dipping_plane
 587            )
 588            * self.cfg.infill_factor
 589            * random_thickness_factor_map
 590        )
 591        thickness_map = np.clip(
 592            thickness_map, 0.0, self.cfg.thickness_max * self.cfg.infill_factor * 1.5
 593        )
 594        return thickness_map
 595
 596    def _create_initial_thickness_factor_map(self):
 597        """
 598        Create the thickness_factor map for the layer at the base of the model
 599
 600        Returns
 601        -------
 602        random_thickness_factor_map : ndarray - a random thickness factor map for the initial layer at base
 603        """
 604        _, random_thickness_factor_map = self._generate_random_depth_structure_map(
 605            dip_range=[0.0, 0.0], num_points_range=(25, 100), elevation_std=0.45
 606        )
 607        random_thickness_factor_map -= random_thickness_factor_map.mean() - 1
 608        #
 609        return random_thickness_factor_map
 610
 611    def create_depth_maps(self):
 612        """Building layers in reverse order - starting at bottom and depositing new layers on top.
 613
 614        Each layer has random residual dip and pseudo-random residual thickness
 615        """
 616        self._generate_lookup_tables()
 617        # Create initial depth map at base of model using initial=True
 618        _, previous_depth_map = self._generate_random_depth_structure_map(
 619            dip_range=[0.0, 75],
 620            num_points_range=(3, 5),
 621            initial=True,
 622            elevation_std=self.cfg.initial_layer_stdev,
 623        )
 624
 625        # Build layers while minimum depth is less than the seabed_min_depth as set in config (given in metres).
 626        # Convert min_depth into samples, don't forget about the infill factor
 627        shallowest_depth_to_build = (
 628            self.cfg.seabed_min_depth / float(self.cfg.digi)
 629        ) * self.cfg.infill_factor
 630
 631        # Build layers in a loop from deep to shallow until the minimum depth is reached and then break out.
 632        for i in range(20000):
 633            # Do the special case for the random_thickness_factor_map for the initial layer (at base)
 634            if i == 0:
 635                if self.cfg.verbose:
 636                    print("Building random depth map at base of model")
 637                random_thickness_factor_map = (
 638                    self._create_initial_thickness_factor_map()
 639                )
 640            else:  # Otherwise create standard random thickness factor map
 641                if self.cfg.verbose:
 642                    print(f"Building Layer {i}")
 643                random_thickness_factor_map = self._random_layer_thickness()
 644            # Create the layer's thickness map using the random_thickness_factor_map
 645            thickness_map = self._create_thickness_map(random_thickness_factor_map, i)
 646            current_depth_map = previous_depth_map - thickness_map
 647
 648            if self.cfg.verbose:
 649                print(
 650                    f"current_depth_map min/mean/max = {current_depth_map.min():.2f}, {current_depth_map.mean():.2f},"
 651                    f" {current_depth_map.max():.2f}"
 652                )
 653                print(
 654                    f"thickness_map min/mean/max = {thickness_map.min():.2f}, {thickness_map.mean():.2f},"
 655                    f" {thickness_map.max():.2f}"
 656                )
 657
 658            # break out of loop when minimum depth is reached
 659            if current_depth_map.min() <= shallowest_depth_to_build:
 660                break
 661
 662            # replace previous depth map for next iteration
 663            previous_depth_map = current_depth_map.copy()
 664            # save depth map in (top of) 3d array. Resulting order is shallow at top
 665            try:
 666                depth_maps = np.dstack((current_depth_map, depth_maps))
 667            except UnboundLocalError:
 668                # There is no stack of horizons yet - write the base.
 669                depth_maps = current_depth_map.copy()
 670            if self.cfg.verbose:
 671                print(f"Layer {i}, depth_maps.shape = {depth_maps.shape}")
 672            self.max_layers = i + 1
 673
 674        if self.cfg.verbose:
 675            print("\n ... finished creating horizon layers ...")
 676        # Store maps in hdf file
 677        self.depth_maps = self.cfg.hdf_init("depth_maps", shape=depth_maps.shape)
 678        self.depth_maps[:] = depth_maps
 679
 680        self.cfg.write_to_logfile(
 681            f"number_layers: {self.max_layers}",
 682            mainkey="model_parameters",
 683            subkey="number_layers",
 684            val=self.max_layers,
 685        )
 686
 687
 688class Onlaps(Horizons):
 689    def __init__(self, parameters, depth_maps, thicknesses, max_layers):
 690        self.cfg = parameters
 691        self.depth_maps = depth_maps
 692        self.thicknesses = thicknesses
 693        self.max_layers = max_layers
 694        self.onlap_horizon_list = list()
 695        self._generate_onlap_lookup_table()
 696
 697    def _generate_onlap_lookup_table(self):
 698        # Onlaps
 699        onlap_layer_list = np.sort(
 700            np.random.uniform(
 701                low=5,
 702                high=self.max_layers,
 703                size=int(np.random.triangular(1, 4, 7) + 0.5),
 704            ).astype("int")
 705        )
 706        if self.cfg.verbose:
 707            print("self.cfg.num_lyr_lut = ", self.cfg.num_lyr_lut)
 708            print("onlap_layer_list = ", onlap_layer_list)
 709        onlap_array_dim = int(500 / 1250 * self.cfg.cube_shape[2])
 710        self.onlaps = np.zeros(onlap_array_dim, "int")
 711        self.onlaps[onlap_layer_list] = 1
 712
 713        if self.cfg.verbose:
 714            print(
 715                f"Number of onlapping flags: {self.onlaps[self.onlaps == 1].shape[0]}"
 716            )
 717
 718    def insert_tilting_episodes(self):
 719        ############################################################################
 720        # insert tilting (onlap) episodes
 721        # - computations loop from deep horizons toward shallow (top of horizon cube)
 722        ############################################################################
 723        if self.cfg.verbose:
 724            print("\n ... simulating tilting (onlap) episodes ...")
 725
 726        azi_list = []
 727        dip_list = []
 728
 729        count = 0
 730        onlaps_horizon_list = []
 731        layer_has_onlaps = np.zeros((self.depth_maps.shape[2]), "bool")
 732        for i in range(self.depth_maps.shape[2] - 2, 0, -1):
 733
 734            if self.onlaps[i] == 1:
 735                # A tilting (onlap) episode occurs for this horizon.
 736                # Add a random dipping layer to this and all shallower layers, and adjust depth for layer thickness
 737                count += 1
 738                if self.cfg.verbose:
 739                    print(
 740                        " ... simulate a tilting (onlap) episode ... at output horizon number {}".format(
 741                            i
 742                        )
 743                    )
 744                onlaps_horizon_list.append(i)
 745
 746                if not self.cfg.regen:
 747                    azi = np.random.uniform(low=0.0, high=360.0)
 748                    azi_list.append(azi)
 749                    dip = np.random.uniform(low=5.0, high=20.0)
 750                    dip_list.append(dip)
 751
 752                dipping_plane2 = (
 753                    self._fit_plane_strike_dip(
 754                        azi_list[count - 1],
 755                        dip_list[count - 1],
 756                        self.depth_maps.shape,
 757                        verbose=False,
 758                    )
 759                    * self.cfg.infill_factor
 760                )
 761                dipping_plane2_offset = (
 762                    self.thicknesses[i] * self.cfg.infill_factor / 2.0
 763                    - dipping_plane2.max()
 764                )
 765
 766                # adjust all shallower layers
 767                previous_depth_map = self.depth_maps[:, :, i + 1].copy()
 768                for i_horizon in range(i, 0, -1):
 769                    current_depth_map = self.depth_maps[:, :, i_horizon].copy()
 770                    thickness_map = previous_depth_map - current_depth_map
 771                    prior_thickness_map = thickness_map.copy()
 772                    thickness_map += dipping_plane2
 773                    thickness_map += dipping_plane2_offset
 774                    thickness_map = np.clip(thickness_map, 0.0, thickness_map.max())
 775                    current_depth_map = previous_depth_map - thickness_map
 776                    if np.mean(thickness_map - prior_thickness_map) > 0.5:
 777                        layer_has_onlaps[i_horizon] = 1
 778                    self.depth_maps[:, :, i_horizon] = current_depth_map.copy()
 779
 780        # print("\n\n ... depth_maps.shape = {}".format(depth_maps.shape))
 781        if self.cfg.verbose:
 782            print(
 783                f" ... Finished inserting tilting (onlap) episodes. {count} episodes were added..."
 784            )
 785        self.cfg.write_to_logfile(
 786            f"number_onlap_episodes: {count}\nonlaps_horizon_list: {str(onlaps_horizon_list)}"
 787        )
 788
 789        self.cfg.write_to_logfile(
 790            msg=None,
 791            mainkey="model_parameters",
 792            subkey="number_onlap_episodes",
 793            val=count,
 794        )
 795        self.cfg.write_to_logfile(
 796            msg=None,
 797            mainkey="model_parameters",
 798            subkey="onlaps_horizon_list",
 799            val=str(onlaps_horizon_list),
 800        )
 801        self.onlap_horizon_list = onlaps_horizon_list
 802        return self.onlap_horizon_list
 803
 804
 805class BasinFloorFans(Horizons):
 806    def __init__(self, parameters, max_layers):
 807        self.cfg = parameters
 808        self.max_layers = max_layers
 809        self.fan_layers = None
 810        # self._generate_fan_lookup_table()
 811
 812    def _generate_fan_lookup_table(self):
 813        # TODO test where & how often fans should be added. Also check if they overlap with onlaps and if so, what to do
 814        # Onlaps
 815        layers_with_fans = np.sort(
 816            np.random.uniform(
 817                low=5, high=self.max_layers - 1, size=int(np.random.choice([1, 2, 3]))
 818            ).astype("int")
 819        )
 820        self.fan_layers = layers_with_fans
 821        self.fan_thicknesses = []
 822        self.fans = np.zeros(self.max_layers, "int")
 823        self.fans[list(self.fan_layers)] = 1
 824
 825    def _generate_basin_floor_fan(self, layer_number):
 826        # Select parameters for fan
 827        _scale = np.random.uniform(50.0, 200.0)  # length in pixels
 828        _aspect = np.random.uniform(1.5, 4.0)  # length  width
 829        _azimuth = np.random.uniform(0.0, 360.0)
 830        _factor = np.random.uniform(1.5, 3.5)
 831        _asymmetry_factor = np.random.uniform(-1.0, 1.0)
 832        _smoothing_size = np.random.uniform(1.33, 3.0)
 833
 834        # Choose whether to create a pair of fans
 835        pair_of_fans = np.random.choice([True, False])
 836
 837        fan_parameters = (
 838            f"{_scale:4.2f}, {_aspect:4.2f}, {_azimuth:4.2f}, {_factor:4.2f}, {_asymmetry_factor:4.2f},"
 839            f" {_smoothing_size:4.2f}"
 840        )
 841        if self.cfg.verbose:
 842            print(fan_parameters)
 843        zi = self._generate_fan_thickness_map(
 844            scale=_scale,
 845            aspect=_aspect,
 846            rotate_angle=_azimuth,
 847            dip=3.5,
 848            entry_point_factor=_factor,
 849            asymmetry_factor=_asymmetry_factor,
 850            smoothing_size=_smoothing_size,
 851        )
 852
 853        if pair_of_fans:
 854            if self.cfg.verbose:
 855                print("Creating a pair of fans\n")
 856            _scale2 = _scale * np.random.uniform(0.667, 1.5)  # length in pixels
 857            _aspect2 = _aspect * np.random.uniform(0.75, 1.33)  # length / width
 858            _azimuth2 = _azimuth + np.random.uniform(-30.0, 30.0)
 859            delta_dist = (_scale + _scale2) / (
 860                (_aspect + _aspect2)
 861                * np.random.triangular(0.5, 0.85, 1.5)
 862                * np.random.choice([-1.0, 1.0])
 863            )  # length in pixels
 864            delta_x, delta_y = self.rotate_point(delta_dist, 0.0, 360.0 - _azimuth2)
 865            _factor2 = _factor + np.random.uniform(0.8, 1.25)
 866            _asymmetry_factor2 = _asymmetry_factor + np.random.uniform(-0.25, 0.25)
 867            _smoothing_size2 = _smoothing_size + np.random.uniform(-0.25, 0.25)
 868
 869            fan_parameters_2 = (
 870                f"{_scale2:4.2f}, {_aspect2:4.2f}, {_azimuth2:4.2f}, {_factor2:4.2f},"
 871                f"{_asymmetry_factor2:4.2f}, {_smoothing_size2:4.2f}"
 872            )
 873
 874            if self.cfg.verbose:
 875                print(
 876                    f"Fan 1 parameters: {fan_parameters}\nFan 2 parameters: {fan_parameters_2}"
 877                )
 878            fan_parameters += f"\n{fan_parameters_2}"
 879            zi2 = self._generate_fan_thickness_map(
 880                scale=_scale2,
 881                aspect=_aspect2,
 882                rotate_angle=_azimuth2,
 883                dip=3.5,
 884                entry_point_factor=_factor2,
 885                asymmetry_factor=_asymmetry_factor2,
 886                smoothing_size=_smoothing_size2,
 887            )
 888
 889            zi2_padded = np.zeros((zi.shape[0] * 2, zi.shape[1] * 2), "float")
 890            zi2_padded[
 891                zi.shape[0] // 2 : zi.shape[0] + zi.shape[0] // 2,
 892                zi.shape[1] // 2 : zi.shape[1] + zi.shape[1] // 2,
 893            ] += zi2
 894            zi2_padded = np.roll(zi2_padded, int(delta_x + 0.5), axis=0)
 895            zi2_padded = np.roll(zi2_padded, int(delta_y + 0.5), axis=1)
 896            zi_delta = (
 897                zi2_padded[
 898                    zi.shape[0] // 2 : zi.shape[0] + zi.shape[0] // 2,
 899                    zi.shape[1] // 2 : zi.shape[1] + zi.shape[1] // 2,
 900                ]
 901                - zi
 902            )
 903            zi_delta[zi_delta < 0.0] = 0.0
 904            zi3 = zi + zi_delta
 905            zi = zi3
 906
 907        # Plot the fan
 908        if self.cfg.qc_plots:
 909            from datagenerator.util import import_matplotlib
 910
 911            plt = import_matplotlib()
 912            plt.figure(1)
 913            plt.clf()
 914            plt.title(fan_parameters)
 915            plt.imshow(zi.T)
 916            plt.colorbar()
 917            plt.savefig(
 918                os.path.join(
 919                    self.cfg.work_subfolder, f"random_fan_layer{layer_number}.png"
 920                )
 921            )
 922
 923        self.fan_thicknesses.append(zi)
 924        return zi
 925
 926    def _generate_fan_thickness_map(
 927        self,
 928        scale=2.5,
 929        aspect=1.5,
 930        rotate_angle=0.0,
 931        zero_at_corners=True,
 932        dip=3.5,
 933        entry_point_factor=1.0,
 934        asymmetry_factor=0.5,
 935        smoothing_size=2.5,
 936    ):
 937        ############################################################################
 938        # generate a 2D array representing a depth (or TWT) structure map.
 939        # - grid_size controls output size in (x,y), units are samples
 940        # - dip_range controls slope of dipping plane upon which random residual points are placed
 941        # - num_points controls number of randomly positioned points are used for residual grid
 942        # - asymetry_factor: range [-1, 1.]. 0. does nothing, sign determines left/right prominence
 943        ############################################################################
 944
 945        grid_size = self.cfg.cube_shape[:2]
 946        # define output grid (padded by 15% in each direction)
 947        xi = np.linspace(-0.15, 1.15, int(float(grid_size[0]) * 1.3))
 948        yi = np.linspace(-0.15, 1.15, int(float(grid_size[1]) * 1.3))
 949
 950        # start with a gently dipping plane similar to that found on passive shelf margins (< 1 degree dip)
 951        dipping_plane = self._fit_plane_strike_dip(
 952            360.0 - rotate_angle, dip, grid_size, verbose=False
 953        ).T
 954
 955        x = np.array(
 956            [
 957                0.5 * (1.0 - np.cos(theta)) * np.sin(theta) * np.cos(phi)
 958                for phi in np.arange(0.0, np.pi, np.pi / 180.0)
 959                for theta in np.arange(0.0, np.pi, np.pi / 180.0)
 960            ]
 961        )
 962        y = np.array(
 963            [
 964                np.cos(theta)
 965                for phi in np.arange(0.0, np.pi, np.pi / 180.0)
 966                for theta in np.arange(0.0, np.pi, np.pi / 180.0)
 967            ]
 968        )
 969        z = np.array(
 970            [
 971                0.5 * (1.0 - np.cos(theta)) * np.sin(theta) * np.sin(phi)
 972                for phi in np.arange(0.0, np.pi, np.pi / 180.0)
 973                for theta in np.arange(0.0, np.pi, np.pi / 180.0)
 974            ]
 975        )
 976        x -= x.min()
 977        y -= y.min()
 978        if self.cfg.verbose:
 979            print("number of xyz points = ", x.size)
 980            print("x min/mean/max = ", x.min(), x.mean(), x.max(), x.max() - x.min())
 981        x *= (scale / aspect) / (x.max() - x.min())
 982        y *= scale / (y.max() - y.min())
 983
 984        # scale back to range of [0,1]
 985        x /= grid_size[0]
 986        y /= grid_size[1]
 987
 988        if self.cfg.verbose:
 989            print("x min/mean/max = ", x.min(), x.mean(), x.max(), x.max() - x.min())
 990            print("y min/mean/max = ", y.min(), y.mean(), y.max(), y.max() - y.min())
 991
 992        # add asymmetry
 993        # - small rotation, stretch, undo rotation
 994        if asymmetry_factor != 0.0:
 995            x, y = self.rotate_point(x, y, asymmetry_factor * 45.0)
 996            x *= 2.0 + np.abs(asymmetry_factor)
 997            y *= 2.0 - np.abs(asymmetry_factor)
 998            x, y = self.rotate_point(x, y, -asymmetry_factor * 45.0)
 999            x /= 2.0
1000            y /= 2.0
1001        if self.cfg.verbose:
1002            print("x min/mean/max = ", x.min(), x.mean(), x.max(), x.max() - x.min())
1003            print("y min/mean/max = ", y.min(), y.mean(), y.max(), y.max() - y.min())
1004
1005        # rotate
1006        x, y = self.rotate_point(x, y, 360.0 - rotate_angle)
1007
1008        # center in the grid
1009        x += 0.5 - x.mean()
1010        y += 0.5 - y.mean()
1011
1012        if self.cfg.verbose:
1013            print("x min/mean/max = ", x.min(), x.mean(), x.max(), x.max() - x.min())
1014            print("y min/mean/max = ", y.min(), y.mean(), y.max(), y.max() - y.min())
1015            print("z min/mean/max = ", z.min(), z.mean(), z.max(), z.max() - z.min())
1016
1017        # decimate the grid randomly
1018        decimation_factor = 250
1019
1020        fan_seed_x = np.random.randint(1, high=2 ** 32 - 1)
1021        fan_seed_y = np.random.randint(1, high=2 ** 32 - 1)
1022        fan_seed_z = np.random.randint(1, high=2 ** 32 - 1)
1023        point_indices = np.arange(x.size).astype("int")
1024        point_indices = point_indices[: x.size // decimation_factor]
1025        np.random.shuffle(point_indices)
1026        np.random.seed(fan_seed_x)
1027        _x = x + np.random.uniform(low=-0.03, high=0.03, size=x.size)
1028        np.random.seed(fan_seed_y)
1029        _y = y + np.random.uniform(low=-0.03, high=0.03, size=x.size)
1030        np.random.seed(fan_seed_z)
1031        _z = z + np.random.uniform(low=-0.1, high=0.1, size=x.size)
1032
1033        _x = _x[point_indices]
1034        _y = _y[point_indices]
1035        _z = _z[point_indices]
1036
1037        if self.cfg.qc_plots:
1038            from datagenerator.util import import_matplotlib
1039
1040            plt = import_matplotlib()
1041            plt.clf()
1042            plt.grid()
1043            plt.scatter(x * 300, y * 300, c=z)
1044            plt.xlim([0, 300])
1045            plt.ylim([0, 300])
1046            plt.colorbar()
1047
1048        if zero_at_corners:
1049            x = np.hstack((x, [-0.15, -0.15, 1.15, 1.15]))
1050            y = np.hstack((y, [-0.15, 1.15, -0.15, 1.15]))
1051            z = np.hstack((z, [-0.15, -0.15, -0.15, -0.15]))
1052            _x = np.hstack((_x, [-0.15, -0.15, 1.15, 1.15]))
1053            _y = np.hstack((_y, [-0.15, 1.15, -0.15, 1.15]))
1054            _z = np.hstack((_z, [-0.15, -0.15, -0.15, -0.15]))
1055
1056        # grid the data.
1057        zi = griddata(
1058            (x, y),
1059            z,
1060            (xi.reshape(xi.shape[0], 1), yi.reshape(1, yi.shape[0])),
1061            method="linear",
1062        )
1063        _zi = griddata(
1064            (_x, _y),
1065            _z,
1066            (xi.reshape(xi.shape[0], 1), yi.reshape(1, yi.shape[0])),
1067            method="cubic",
1068        )
1069
1070        zi[np.isnan(zi)] = 0.0
1071        zi[zi <= 0.0] = 0.0
1072        zi_mask = zi.copy()
1073        zi_mask[zi_mask <= 0.015] = 0.0
1074        if self.cfg.verbose:
1075            print("zi_mask[zi_mask >0.].min() = ", zi_mask[zi_mask > 0.0].min())
1076            print("zi_mask[zi_mask >0.].mean() = ", zi_mask[zi_mask > 0.0].mean())
1077            print(
1078                "np.percentile(zi_mask[zi_mask >0.],.5) = ",
1079                np.percentile(zi_mask[zi_mask > 0.0], 0.5),
1080            )
1081
1082        _zi[np.isnan(_zi)] = 0.0
1083        zi = _zi + 0.0
1084        zi[zi_mask <= 0.0] = 0.0
1085        zi[zi <= 0.0] = 0.0
1086        from scipy.ndimage import gaussian_filter, grey_closing
1087
1088        zi = gaussian_filter(zi, smoothing_size)
1089        # For cube_shape 300, 300, use closing size of (10, 10)
1090        closing_size = (
1091            int(self.cfg.cube_shape[0] / 30),
1092            int(self.cfg.cube_shape[1] / 30),
1093        )
1094        zi = grey_closing(zi, size=closing_size)
1095
1096        zi = zi.reshape(xi.shape[0], yi.shape[0])
1097        xi_min_index = np.argmin((xi - 0.0) ** 2)
1098        xi_max_index = xi_min_index + grid_size[0]
1099        yi_min_index = np.argmin((yi - 0.0) ** 2)
1100        yi_max_index = yi_min_index + grid_size[1]
1101        zi = zi[xi_min_index:xi_max_index, yi_min_index:yi_max_index]
1102        # add to the gently sloping plane
1103        dipping_plane[zi == 0.0] = 0.0
1104        dipping_plane -= dipping_plane[dipping_plane != 0.0].min()
1105        dipping_plane[zi <= 0.05] = 0.0
1106        if dipping_plane.max() > 0:
1107            dipping_plane /= dipping_plane[dipping_plane != 0.0].max()
1108        dipping_plane *= zi.max()
1109        if self.cfg.verbose:
1110            print(
1111                "zi[zi > 0.] min/mean/max = ",
1112                zi[zi > 0.0].min(),
1113                zi[zi > 0.0].mean(),
1114                zi[zi > 0.0].max(),
1115            )
1116        if dipping_plane.max() > 0:
1117            if self.cfg.verbose:
1118                print(
1119                    "dipping_plane[dipping_plane > 0.] min/mean/max = ",
1120                    dipping_plane[dipping_plane > 0.0].min(),
1121                    dipping_plane[dipping_plane > 0.0].mean(),
1122                    dipping_plane[dipping_plane > 0.0].max(),
1123                )
1124        else:
1125            if self.cfg.verbose:
1126                print(
1127                    "dipping_plane min/mean/max = ",
1128                    dipping_plane.min(),
1129                    dipping_plane.mean(),
1130                    dipping_plane.max(),
1131                )
1132        zi += entry_point_factor * dipping_plane
1133
1134        # todo QC size of this smoothing (to smooth edges of fan)
1135        zi = gaussian_filter(zi, smoothing_size * 2.0)
1136
1137        return zi
1138
1139    def insert_fans_into_horizons(self, depth_maps):
1140        new_depth_maps = depth_maps[:].copy()
1141        # Create lookup table for fans
1142        self._generate_fan_lookup_table()
1143        # generate fans using fan_layers
1144        for layer in self.fan_layers:
1145            thickness_map = (
1146                self._generate_basin_floor_fan(layer) * self.cfg.infill_factor
1147            )
1148            new_depth_maps = self.insert_feature_into_horizon_stack(
1149                thickness_map, layer, new_depth_maps
1150            )
1151        # Write fan layers to logfile
1152        self.cfg.write_to_logfile(
1153            f"number_fan_episodes: {len(self.fan_layers)}\n"
1154            f"fans_horizon_list: {str(self.fan_layers)}"
1155        )
1156        self.cfg.write_to_logfile(
1157            msg=None,
1158            mainkey="model_parameters",
1159            subkey="number_fan_episodes",
1160            val=len(self.fan_layers),
1161        )
1162        self.cfg.write_to_logfile(
1163            msg=None,
1164            mainkey="model_parameters",
1165            subkey="fan_horizon_list",
1166            val=str(self.fan_layers),
1167        )
1168
1169        return new_depth_maps
1170
1171    def fan_qc_plot(self, maps, lyrnum, thickness):
1172        """
1173        Plot a cross-section of the inserted basin floor fan
1174
1175        Display an inline and crossline with the basin floor fan in red in cross-section
1176        Inlay a map of the fan in each subplot
1177        """
1178        from datagenerator.util import import_matplotlib
1179
1180        plt = import_matplotlib()
1181        plt.close()
1182        fig, axs = plt.subplots(nrows=2, ncols=1, figsize=(20, 15), sharey=True)
1183        axs[0].set_title(f"Layer {lyrnum}")
1184        axs[1].set_title(f"Layer {lyrnum}")
1185        #
1186        max_layer = np.clip(
1187            lyrnum + 10, 0, maps.shape[2]
1188        )  # only plot layers below fan if they exist
1189        for j in range(lyrnum - 5, max_layer, 1):
1190            if j == lyrnum:  # highlight fan layers in thicker red line
1191                axs[0].plot(
1192                    range(self.cfg.cube_shape[0]),
1193                    maps[int(self.cfg.cube_shape[0] / 2), :, j],
1194                    "r-",
1195                    lw=0.5,
1196                )
1197                axs[1].plot(
1198                    range(self.cfg.cube_shape[0]),
1199                    maps[:, int(self.cfg.cube_shape[1] / 2), j],
1200                    "r-",
1201                    lw=0.5,
1202                )
1203            else:
1204                axs[0].plot(
1205                    range(self.cfg.cube_shape[0]),
1206                    maps[int(self.cfg.cube_shape[0] / 2), :, j],
1207                    "k-",
1208                    lw=0.2,
1209                )
1210                axs[1].plot(
1211                    range(self.cfg.cube_shape[0]),
1212                    maps[:, int(self.cfg.cube_shape[1] / 2), j],
1213                    "k-",
1214                    lw=0.2,
1215                )
1216        for inl, ax in enumerate(axs):
1217            if inl == 0:
1218                plt.axes([0.8, 0.75, 0.11, 0.11])
1219                plt.axvline(
1220                    x=int(self.cfg.cube_shape[0] / 2), color="grey", linestyle="--"
1221                )
1222            else:
1223                plt.axes([0.8, 0.33, 0.11, 0.11])
1224                plt.axhline(
1225                    y=int(self.cfg.cube_shape[1] / 2), color="grey", linestyle="--"
1226                )
1227            plt.imshow(thickness.T)
1228            plt.xticks([])
1229            plt.yticks([])
1230            ax.set_ylim(
1231                top=np.min(maps[..., lyrnum - 5]),
1232                bottom=np.max(maps[..., max_layer - 1]),
1233            )  # Reverse Y axis
1234        fig.savefig(
1235            os.path.join(self.cfg.work_subfolder, f"BasinFloorFan_layer_{lyrnum}.png")
1236        )
1237        plt.close()
1238
1239
1240class Channel(Horizons):
1241    def __init__(self):
1242        self.W = 200.0  # Channel wodth
1243        self.D = 12.0  # Channel depth
1244        self.pad = 100  # padding (number of nodepoints along centerline)
1245        self.deltas = 50.0  # sampling distance along centerline
1246        self.nit = 2000  # number of iterations
1247        self.Cf = 0.022  # dimensionless Chezy friction factor
1248        self.crdist = 1.5 * self.W  # threshold distance at which cutoffs occur
1249        self.kl = 60.0 / (365 * 24 * 60 * 60.0)  # migration rate constant (m/s)
1250        self.kv = 1.0e-11  # vertical slope-dependent erosion rate constant (m/s)
1251        self.dt = 2 * 0.05 * 365 * 24 * 60 * 60.0  # time step (s)
1252        self.dens = 1000  # density of water (kg/m3)
1253        self.saved_ts = 20  # which time steps will be saved
1254        self.n_bends = 30  # approximate number of bends to model
1255        self.Sl = 0.0  # initial slope (matters more for submarine channels than rivers)
1256        self.t1 = 500  # time step when incision starts
1257        self.t2 = 700  # time step when lateral migration starts
1258        self.t3 = 1400  # time step when aggradation starts
1259        self.aggr_factor = (
1260            4e-9  # aggradation factor (m/s, about 0.18 m/year, it kicks in after t3)
1261        )
1262
1263        self.h_mud = 0.3  # thickness of overbank deposits for each time step
1264        self.dx = 20.0  # gridcell size in metres
1265
1266        # Channel objects
1267        self.ch = None
1268        self.chb = None
1269        self.chb_3d = None
1270
1271    def generate_channel_parameters(self):
1272        # select random parameters
1273        pass
1274
1275    def create_channel_3d(self):
1276        # Initialise channel
1277        self.ch = mp.generate_initial_channel(
1278            self.W, self.D, self.Sl, self.deltas, self.pad, self.n_bends
1279        )
1280        # Create channel belt object
1281        self.chb = mp.ChannelBelt(
1282            channels=[self.ch], cutoffs=[], cl_times=[0.0], cutoff_times=[]
1283        )
1284        # Migrate channel
1285        self.chb.migrate(
1286            self.nit,
1287            self.saved_ts,
1288            self.deltas,
1289            self.pad,
1290            self.crdist,
1291            self.Cf,
1292            self.kl,
1293            self.kv,
1294            self.dt,
1295            self.dens,
1296            self.t1,
1297            self.t2,
1298            self.t3,
1299            self.aggr_factor,
1300        )
1301        end_time = self.chb.cl_times[-1]
1302        _xmin = 15000
1303        _xmax = 21000
1304        self.chb_3d, _, _, _, _ = self.chb.build_3d_model(
1305            "fluvial",
1306            h_mud=self.h_mud,
1307            levee_width=800.0,
1308            h=12.0,
1309            w=self.W,
1310            bth=0.0,
1311            dcr=10.0,
1312            dx=self.dx,
1313            delta_s=self.deltas,
1314            starttime=self.chb.cl_times[0],
1315            endtime=end_time,
1316            xmin=_xmin,
1317            xmax=_xmax,
1318            ymin=-3000,
1319            ymax=3000,
1320        )
1321
1322    def insert_channel_into_horizons(self, layer, depth_maps):
1323        new_depth_maps = self.insert_feature_into_horizon_stack(
1324            self.chb.strat, layer, depth_maps
1325        )
1326        return new_depth_maps
1327
1328    def insert_channel_facies(self):
1329        pass
1330
1331
1332class Facies:
1333    def __init__(self, parameters, max_layers, onlap_horizon_list, fan_horizon_list):
1334        self.cfg = parameters
1335        self.max_layers = max_layers
1336        self.onlap_horizon_list = onlap_horizon_list
1337        self.fan_horizon_list = fan_horizon_list
1338        self.facies = None
1339
1340    def sand_shale_facies_binomial_dist(self):
1341        """Randomly select sand or shale facies using binomial distribution using the sand_layer_pct from config file"""
1342        sand_layer = np.random.binomial(
1343            1, self.cfg.sand_layer_pct, size=self.max_layers
1344        )
1345        # Insert a water-layer code of -1 at the top
1346        self.facies = np.hstack((np.array((-1.0)), sand_layer))
1347
1348    def sand_shale_facies_markov(self):
1349        """Generate a 1D array of facies usinga 2-state Markov process
1350
1351        Note: the Binomial distribution can be generated by setting the sand_layer_thickness to 1/sand_layer_pct
1352        """
1353        mk = MarkovChainFacies(
1354            self.cfg.sand_layer_pct, self.cfg.sand_layer_thickness, (0, 1)
1355        )
1356        # Randomly select initial state
1357        facies = mk.generate_states(np.random.choice(2, 1)[0], num=self.max_layers)
1358        self.facies = np.hstack((np.array((-1.0)), facies))
1359
1360    def set_layers_below_onlaps_to_shale(self):
1361        """
1362        Set layers immediately below tilting episode (onlap surface) to shale.
1363
1364        Set facies array to 0 (shale) above any layer which is marked as being an onlap surface
1365        """
1366        onlap_horizon_list_plus_one = np.array(self.onlap_horizon_list)
1367        onlap_horizon_list_plus_one += 1
1368        self.facies[onlap_horizon_list_plus_one] = 0.0
1369
1370    def set_fan_facies(self, depth_maps):
1371        """Set fan facies as sand layers surrounded by shales
1372
1373        Set the fan layer to be a sand.
1374        Set the layers above and below the fan to be shale
1375        """
1376        # Find those layers which are directly above (and touching) the fan layer
1377        layers_with_zero_thickness = set(np.where(np.diff(depth_maps) == 0)[2])
1378        layers_above_fan = []
1379        for f in self.fan_horizon_list:
1380            for i in range(f - 1, 1, -1):
1381                if i in layers_with_zero_thickness:
1382                    layers_above_fan.append(i)
1383                else:
1384                    break
1385        sand = list(self.fan_horizon_list)
1386        shale = layers_above_fan + list(self.fan_horizon_list + 1)
1387        self.facies[sand] = 1.0
1388        self.facies[shale] = 0.0
1389
1390
1391class MarkovChainFacies:
1392    def __init__(self, sand_fraction, sand_thickness, states=(0, 1)):
1393        """Initialize the MarkovChain instance.
1394
1395        Parameters
1396        ----------
1397        sand_fraction : float
1398            Fraction of sand layers in model
1399        sand_thickness : int
1400            Thickness of sand layers, given in units of layers
1401        states : iterable
1402            An iterable representing the states of the Markov Chain.
1403        """
1404        self.sand_fraction = sand_fraction
1405        self.sand_thickness = sand_thickness
1406        self.states = states
1407        self.index_dict = {
1408            self.states[index]: index for index in range(len(self.states))
1409        }
1410        self.state_dict = {
1411            index: self.states[index] for index in range(len(self.states))
1412        }
1413        self._transition_matrix()
1414
1415    def _transition_matrix(self):
1416        """TODO Stephs notes go here"""
1417        beta = 1 / self.sand_thickness
1418        alpha = self.sand_fraction / (self.sand_thickness * (1 - self.sand_fraction))
1419        self.transition = np.array([[1 - alpha, alpha], [beta, 1 - beta]])
1420
1421    def next_state(self, current_state):
1422        """Returns the state of the random variable at the next instance.
1423
1424        Parameters
1425        ----------
1426        current_state :str
1427            The current state of the system
1428        """
1429        return np.random.choice(
1430            self.states, p=self.transition[self.index_dict[current_state], :]
1431        )
1432
1433    def generate_states(self, current_state, num=100):
1434        """Generate states of the system with length num
1435
1436        Parameters
1437        ----------
1438        current_state : str
1439            The state of the current random variable
1440        num : int, optional
1441            [description], by default 100
1442        """
1443        future_states = []
1444        for _ in range(num):
1445            next_state = self.next_state(current_state)
1446            future_states.append(next_state)
1447            current_state = next_state
1448        return np.array(future_states)
1449
1450
1451def build_unfaulted_depth_maps(parameters: Parameters):
1452    """
1453        Build Unfaulted Depth Maps
1454        --------------------------
1455        Generates unfaulted depth maps.
1456
1457        1. Build stack of horizons.
1458        2. Generate a random stack of horizons.
1459        3. Optionally insert basin floor fans
1460        4. Insert onlap episodes
1461
1462        Parameters
1463        ----------
1464        parameters : str
1465            The key desired to be accessed
1466
1467        Returns
1468        -------
1469        depth_maps : np.array
1470            The generated depth maps
1471        onlap_horizon_list : list
1472            Onlapping Horizon list
1473        fan_list : np.array | None
1474            List of generated fans
1475        fan_thicknesses : np.array | None
1476            Generated fan thicknesses
1477    """
1478    horizons = RandomHorizonStack(parameters)
1479    horizons.create_depth_maps()
1480    # Insert onlap episodes
1481    onlaps = Onlaps(
1482        parameters, horizons.depth_maps, horizons.thicknesses, horizons.max_layers
1483    )
1484    onlap_horizon_list = onlaps.insert_tilting_episodes()
1485    # Insert seafloor
1486    depth_maps = horizons.insert_seafloor(horizons.depth_maps)
1487
1488    fan_list = None
1489    fan_thicknesses = None
1490    if parameters.basin_floor_fans:
1491        bff = BasinFloorFans(parameters, horizons.max_layers)
1492        # Insert Fans
1493        depth_maps = bff.insert_fans_into_horizons(horizons.depth_maps)
1494        for layer, thickness in zip(bff.fan_layers, bff.fan_thicknesses):
1495            bff.fan_qc_plot(depth_maps, layer, thickness)
1496        fan_list = bff.fan_layers
1497        fan_thicknesses = bff.fan_thicknesses
1498    return depth_maps, onlap_horizon_list, fan_list, fan_thicknesses
1499
1500
1501def create_facies_array(
1502    parameters: Parameters,
1503    depth_maps: np.ndarray,
1504    onlap_horizons: list,
1505    fan_horizons: np.ndarray = None
1506) -> np.ndarray:
1507    """
1508        Create Facies Array
1509        --------------------------
1510        Generates facies for the model and return an array with the facies.
1511
1512        Parameters
1513        ----------
1514        parameters : datagenerator.Parameters
1515            Parameter object storing all model parameters.
1516        depth_maps : np.ndarray
1517            A numpy array containing the depth maps.
1518        onlap_horizons : list
1519            A list of the onlap horizons.
1520        fan_horizons : np.ndarray
1521            The fan horizons.
1522
1523        Returns
1524        -------
1525        facies : np.ndarray
1526            An array that contains the facies for the model.
1527    """
1528    facies = Facies(parameters, depth_maps.shape[-1], onlap_horizons, fan_horizons)
1529    # facies.sand_shale_facies_binomial_dist()
1530    facies.sand_shale_facies_markov()
1531    if len(onlap_horizons) > 0:
1532        facies.set_layers_below_onlaps_to_shale()
1533    if np.any(fan_horizons):
1534        facies.set_fan_facies(depth_maps)
1535    # Write sand layer % to logfile
1536    sand_pct = facies.facies[facies.facies == 1].size / facies.facies.size
1537    parameters.write_to_logfile(
1538        f"Percent of layers containing sand a priori: {parameters.sand_layer_pct:0.2%}",
1539        mainkey="model_parameters",
1540        subkey="sand_layer_thickness_a_priori",
1541        val=parameters.sand_layer_pct,
1542    )
1543    parameters.write_to_logfile(
1544        f"Sand unit thickness (in layers) a priori: {parameters.sand_layer_thickness}",
1545        mainkey="model_parameters",
1546        subkey="sand_unit_thickness_a_priori",
1547        val=parameters.sand_layer_thickness,
1548    )
1549    parameters.write_to_logfile(
1550        f"Percent of layers containing sand a posteriori: {sand_pct:.2%}",
1551        mainkey="model_parameters",
1552        subkey="sand_layer_percent_a_posteriori",
1553        val=sand_pct,
1554    )
1555    # Count consecutive occurrences (i.e. facies thickness in layers)
1556    facies_groups = [(k, sum(1 for i in g)) for k, g in groupby(facies.facies)]
1557    if parameters.verbose:
1558        print(f"Facies units: {facies_groups}")
1559    avg_sand_thickness = np.nan_to_num(
1560        np.mean([b for a, b in facies_groups if a == 1.0])
1561    )  # convert nan to 0
1562    parameters.write_to_logfile(
1563        f"Average Sand unit thickness (in layers) a posteriori: {avg_sand_thickness:.1f}",
1564        mainkey="model_parameters",
1565        subkey="sand_unit_thickness_a_posteriori",
1566        val=avg_sand_thickness,
1567    )
1568
1569    return facies.facies
1570
1571
1572if __name__ == "__main__":
1573    par = Parameters(user_config="../config/config_bps.json", test_mode=100)
1574    par.setup_model()
1575    # p.cube_shape = (100, 100, 1250)
1576    zmaps, onlaps, fan_list, fan_thickness = build_unfaulted_depth_maps(par)
1577    facies = create_facies_array(par, zmaps, onlaps, fan_list)
1578    print(facies)