In areas with faults and poor signal/noise ratio, where reflectors can be discontinuous from place to place, a dip-based flattening technique might not be able to appropriately track sedimentary layers. To aid the flattening algorithms, one or few reflectors can be picked. This information can be then incorporated in our algorithms as geological constraints. In a first method, we add a model mask to a time domain solution using a Gauss-Newton approach that incorporates an initial solution. In a second method, we set the lower and upper bounds of a constrained optimization algorithm called limited memory BFGS with bounds (L-BFGS-B). Having incorporated the geological information, the flattening algorithms can accurately pick reflectors in 2D and 3D for noisy field data examples. In addition, preliminary results seem to indicate that the L-BFGS-B method converges faster than the Gauss-Newton method.